parser.py

"""
Parses policy text into sections and maps them to categories.
"""

import re

from rules import CATEGORIES


def split_into_sections(text):
    """Split policy text into logical sections based on headings."""
    sections = []

    # Split on heading patterns
    heading_pattern = re.compile(
        r"(?:^|\n)(?:"
        r"#{1,3}\s+(.+)"                           # Markdown headings
        r"|([A-Z][A-Za-z0-9\s&,:'\-]{2,60})\n"    # All-caps or Title Case lines
        r"|(\d+[\.\)]\s*[A-Z].{2,60})"             # Numbered headings
        r")",
        re.MULTILINE,
    )

    lines = text.split("\n")
    current_heading = "Introduction"
    current_content = []

    for line in lines:
        stripped = line.strip()
        if not stripped:
            current_content.append("")
            continue

        # Check if this line looks like a heading
        is_heading = False

        # Markdown heading
        md_match = re.match(r"^#{1,6}\s+(.+)$", stripped)
        if md_match:
            is_heading = True
            new_heading = md_match.group(1).strip()
        # All-caps line (short enough to be a heading)
        elif (
            stripped.isupper()
            and len(stripped) > 3
            and len(stripped) < 80
            and not stripped.endswith(".")
        ):
            is_heading = True
            new_heading = stripped.title()
        # Numbered heading
        elif re.match(r"^[\d\.\)]+\s+[A-Z].{2,60}$", stripped):
            is_heading = True
            new_heading = stripped
        # Bold text that looks like a heading (if surrounded by whitespace)
        elif re.match(r"^\*\*(.+)\*\*$", stripped):
            is_heading = True
            new_heading = re.sub(r"\*\*", "", stripped)
        # Line ending with colon that's short and starts with capital
        elif (
            stripped.endswith(":")
            and len(stripped) < 80
            and stripped[0].isupper()
            and len(stripped.split()) <= 10
            and not any(c.isdigit() for c in stripped[:-1])
        ):
            is_heading = True
            new_heading = stripped.rstrip(":")

        if is_heading:
            # Save previous section
            content = "\n".join(current_content).strip()
            if content or current_heading != "Introduction":
                sections.append({
                    "heading": current_heading,
                    "content": content,
                    "word_count": len(content.split()),
                })
            current_heading = new_heading
            current_content = []
        else:
            current_content.append(stripped)

    # Save last section
    content = "\n".join(current_content).strip()
    sections.append({
        "heading": current_heading,
        "content": content,
        "word_count": len(content.split()),
    })

    return [s for s in sections if s["content"]]


def categorize_sections(sections):
    """Map each section to a privacy policy category based on keywords."""
    categorized = {}

    for cat_key, cat_info in CATEGORIES.items():
        categorized[cat_key] = {
            "name": cat_info["name"],
            "description": cat_info["description"],
            "sections": [],
        }

    for section in sections:
        text_lower = (section["heading"] + " " + section["content"]).lower()
        best_match = None
        best_score = 0

        for cat_key, cat_info in CATEGORIES.items():
            score = 0
            for keyword in cat_info["keywords"]:
                if keyword.lower() in text_lower:
                    score += 1

            if score > best_score:
                best_score = score
                best_match = cat_key

        if best_match and best_score > 0:
            categorized[best_match]["sections"].append(section)
        else:
            # Assign to "other" or first section as general
            if "other" not in categorized:
                categorized["other"] = {
                    "name": "Other / Unclassified",
                    "description": "Content that doesn't fit standard categories",
                    "sections": [],
                }
            categorized["other"]["sections"].append(section)

    return categorized


def get_full_text(sections):
    """Get the full text from all sections."""
    return "\n\n".join(
        f"## {s['heading']}\n{s['content']}" for s in sections
    )


def get_word_count(sections):
    """Total word count across all sections."""
    return sum(s["word_count"] for s in sections)