app/engine/normalize/__init__.py

"""Cleaning, punctuation normalization, and block tagging."""

from __future__ import annotations

import re

from app.engine.models import DocumentBlock

_BIBLIO_HEADINGS = re.compile(
    r"^(#{1,6}\s*)?(references|bibliography|works\s+cited|citations)\s*:?\s*$",
    re.I | re.M,
)
_HEADING = re.compile(r"^(#{1,6}\s+\S.*|[A-Z][A-Z0-9 ,.'-]{2,60})$")
_LIST_ITEM = re.compile(r"^(\d+[\.\)]\s+|[-*•]\s+)\S")
_CODE_FENCE = re.compile(r"^```")
_TABLE_LINE = re.compile(r"^\s*\|.+\|\s*$")
_TABLE_SEPARATOR = re.compile(
    r"^\s*\|?\s*:?-{3,}:?\s*(?:\|\s*:?-{3,}:?\s*)+\|?\s*$"
)
_HTML_TABLE = re.compile(r"</?table\b", re.I)
_FORMULA = re.compile(r"(\$\$.+?\$\$|\\\[[\s\S]+?\\\]|\\begin\{(?:equation|align|math)\})")
_FORMULA_BLOCK = re.compile(
    r"(\$\$|\\\[|\\\]|\\begin\{(?:equation|align|math)\}|"
    r"\\end\{(?:equation|align|math)\})"
)
_URL = re.compile(r"https?://[^\s<>\"']+|www\.[^\s<>\"']+", re.I)
_EMAIL = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b")
_SOFT_WRAP = re.compile(r"(?<=\w)\n(?=\w)")


def normalize_text(text: str) -> str:
    """UTF-8 hygiene, quotes/dashes, soft wraps, spacing."""
    t = (text or "").replace("\r\n", "\n").replace("\r", "\n")
    # Smart quotes / dashes → ASCII equivalents
    t = (
        t.replace("\u201c", '"')
        .replace("\u201d", '"')
        .replace("\u2018", "'")
        .replace("\u2019", "'")
        .replace("\u2013", "-")
        .replace("\u2014", "—")
        .replace("\u00a0", " ")
    )
    # Fix broken line wraps inside paragraphs (keep blank-line paragraph breaks)
    parts = re.split(r"(\n\s*\n)", t)
    fixed: list[str] = []
    for part in parts:
        if re.match(r"\n\s*\n", part):
            fixed.append("\n\n")
            continue
        # Don't join wraps inside code or block-formula fences.
        if "```" in part or _FORMULA_BLOCK.search(part):
            fixed.append(part)
            continue
        lines = part.split("\n")
        rebuilt: list[str] = []
        buf = ""
        for line in lines:
            stripped = line.strip()
            if not stripped:
                if buf:
                    rebuilt.append(buf)
                    buf = ""
                rebuilt.append("")
                continue
            if _LIST_ITEM.match(stripped) or _HEADING.match(stripped) or _TABLE_LINE.match(stripped):
                if buf:
                    rebuilt.append(buf)
                    buf = ""
                rebuilt.append(stripped)
                continue
            if buf and not buf.endswith(("-", "—")):
                # Soft wrap: join with space
                buf = f"{buf} {stripped}"
            elif buf and buf.endswith(("-", "—")):
                buf = buf.rstrip("-—") + stripped
            else:
                buf = stripped
        if buf:
            rebuilt.append(buf)
        fixed.append("\n".join(rebuilt))
    t = "".join(fixed)
    t = re.sub(r"[ \t]+", " ", t)
    t = re.sub(r" *\n *", "\n", t)
    t = re.sub(r"\n{3,}", "\n\n", t)
    # Punctuation spacing
    t = re.sub(r"\s+([,.;:!?])", r"\1", t)
    t = re.sub(r"([(\[{])\s+", r"\1", t)
    t = re.sub(r"\s+([)\]}])", r"\1", t)
    return t.strip()


def _is_bibliography_heading(line: str) -> bool:
    return bool(_BIBLIO_HEADINGS.match(line.strip()))


def _is_markdown_table_block(text: str) -> bool:
    lines = [line.strip() for line in (text or "").splitlines() if line.strip()]
    if len(lines) < 2:
        return False
    table_lines = sum(bool(_TABLE_LINE.match(line)) for line in lines)
    return table_lines == len(lines) and any(
        _TABLE_SEPARATOR.match(line) for line in lines
    )


def detect_blocks(text: str) -> list[DocumentBlock]:
    """Split normalized text into typed blocks with rewrite flags."""
    raw = text or ""
    if not raw.strip():
        return []

    blocks: list[DocumentBlock] = []
    chunks = re.split(r"(\n\s*\n)", raw)
    in_code = False
    in_biblio = False
    buf_lines: list[str] = []
    buf_kind = "paragraph"
    index = 0

    def flush() -> None:
        nonlocal index, buf_lines, buf_kind
        if not buf_lines:
            return
        body = "\n".join(buf_lines).strip("\n")
        if body.strip() == "" and body:
            blocks.append(
                DocumentBlock(text=body, kind="blank", rewriteable=False, index=index)
            )
        elif body.strip():
            rewriteable = buf_kind == "paragraph" and not in_biblio
            blocks.append(
                DocumentBlock(
                    text=body.strip(),
                    kind=buf_kind if not in_biblio or buf_kind != "paragraph" else "bibliography",
                    rewriteable=rewriteable and buf_kind == "paragraph",
                    index=index,
                )
            )
        index += 1
        buf_lines = []
        buf_kind = "paragraph"

    for chunk in chunks:
        if re.match(r"\n\s*\n", chunk or ""):
            flush()
            blocks.append(
                DocumentBlock(text="", kind="blank", rewriteable=False, index=index)
            )
            index += 1
            continue

        stripped_chunk = (chunk or "").strip()
        if (
            _is_markdown_table_block(stripped_chunk)
            or _HTML_TABLE.search(stripped_chunk)
        ):
            flush()
            blocks.append(
                DocumentBlock(
                    text=stripped_chunk,
                    kind="table",
                    rewriteable=False,
                    index=index,
                )
            )
            index += 1
            continue
        if _FORMULA_BLOCK.search(stripped_chunk):
            flush()
            blocks.append(
                DocumentBlock(
                    text=stripped_chunk,
                    kind="formula",
                    rewriteable=False,
                    index=index,
                )
            )
            index += 1
            continue

        for line in (chunk or "").split("\n"):
            stripped = line.strip()

            if _CODE_FENCE.match(stripped):
                if in_code:
                    buf_lines.append(line)
                    buf_kind = "code"
                    flush()
                    in_code = False
                else:
                    flush()
                    in_code = True
                    buf_kind = "code"
                    buf_lines = [line]
                continue

            if in_code:
                buf_lines.append(line)
                buf_kind = "code"
                continue

            if _is_bibliography_heading(stripped):
                flush()
                in_biblio = True
                buf_kind = "bibliography"
                buf_lines = [stripped]
                flush()
                continue

            if (
                in_biblio
                and _HEADING.match(stripped)
                and len(stripped.split()) <= 12
                and not stripped.endswith((".", ",", ";"))
            ):
                flush()
                in_biblio = False
                blocks.append(
                    DocumentBlock(
                        text=stripped,
                        kind="heading",
                        rewriteable=False,
                        index=index,
                    )
                )
                index += 1
                continue

            if _HTML_TABLE.search(stripped) or _TABLE_LINE.match(stripped):
                flush()
                blocks.append(
                    DocumentBlock(
                        text=stripped, kind="table", rewriteable=False, index=index
                    )
                )
                index += 1
                continue

            if _FORMULA.search(stripped) and len(stripped.split()) < 40:
                flush()
                blocks.append(
                    DocumentBlock(
                        text=stripped, kind="formula", rewriteable=False, index=index
                    )
                )
                index += 1
                continue

            if _LIST_ITEM.match(stripped):
                flush()
                blocks.append(
                    DocumentBlock(
                        text=stripped, kind="list", rewriteable=False, index=index
                    )
                )
                index += 1
                continue

            if _HEADING.match(stripped) and len(stripped.split()) <= 12:
                flush()
                blocks.append(
                    DocumentBlock(
                        text=stripped, kind="heading", rewriteable=False, index=index
                    )
                )
                index += 1
                continue

            if in_biblio:
                buf_kind = "bibliography"
            else:
                buf_kind = "paragraph"
            buf_lines.append(stripped if stripped else line)

    flush()

    # Protect paragraphs that are mostly URL/email-only
    for b in blocks:
        if not b.rewriteable:
            continue
        plain = _URL.sub("", b.text)
        plain = _EMAIL.sub("", plain).strip()
        if len(plain.split()) < 3 and (_URL.search(b.text) or _EMAIL.search(b.text)):
            b.rewriteable = False
            b.kind = "special"

    return blocks


def mask_protected_spans(text: str) -> tuple[str, dict[str, str]]:
    """Replace URLs/emails/paths with placeholders for rewrite safety."""
    mapping: dict[str, str] = {}
    counter = {"n": 0}

    def _sub(pattern: re.Pattern[str], label: str, s: str) -> str:
        def repl(m: re.Match[str]) -> str:
            # Alphanumeric placeholders stay as one token in spaCy. Underscore
            # placeholders can be split and lose their trailing delimiter.
            key = f"ZZPROTECTED{label}{counter['n']}ZZ"
            mapping[key] = m.group(0)
            counter["n"] += 1
            return key

        return pattern.sub(repl, s)

    out = text
    out = _sub(_URL, "URL", out)
    out = _sub(_EMAIL, "EMAIL", out)

    def path_repl(m: re.Match[str]) -> str:
        key = f"ZZPROTECTEDPATH{counter['n']}ZZ"
        mapping[key] = m.group(0)
        counter["n"] += 1
        return key

    out = re.sub(r"(?:[A-Za-z]:\\|/)[^\s<>\"']+", path_repl, out)
    return out, mapping


def unmask_protected_spans(text: str, mapping: dict[str, str]) -> str:
    out = text
    for key, val in mapping.items():
        out = out.replace(key, val)
    return out