"""LaTeX parsing utilities.

Provides helpers to extract sections, chunk text, and convert LaTeX to plain
readable prose by stripping markup that is not useful for NLP/embedding.
"""

from __future__ import annotations

import re

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

_SECTION_RE = re.compile(
    r"\\(?:sub)*section\*?\s*\{([^}]*)\}",
    re.DOTALL,
)

# Patterns for LaTeX noise removal (applied in order).
_COMMENT_RE = re.compile(r"%[^\n]*")
_CITE_RE = re.compile(r"\\cite\*?\{[^}]*\}")
_LABEL_RE = re.compile(r"\\label\{[^}]*\}")
_REF_RE = re.compile(r"\\ref\{[^}]*\}")
# Display math: $$...$$  (before single $ to avoid greedy mismatch)
_DISPLAY_DOLLAR_RE = re.compile(r"\$\$.*?\$\$", re.DOTALL)
# Inline math: $...$
_INLINE_MATH_RE = re.compile(r"\$[^$\n]*?\$", re.DOTALL)
# \begin{equation}...\end{equation}
_ENV_EQUATION_RE = re.compile(
    r"\\begin\{equation\*?\}.*?\\end\{equation\*?\}",
    re.DOTALL,
)
# \begin{align}...\end{align}  (covers align, align*, aligned, …)
_ENV_ALIGN_RE = re.compile(
    r"\\begin\{align[^}]*\}.*?\\end\{align[^}]*\}",
    re.DOTALL,
)
# Collapse excess whitespace
_WHITESPACE_RE = re.compile(r"[ \t]+")
_BLANK_LINES_RE = re.compile(r"\n{3,}")


def _clean_latex(text: str) -> str:
    """Strip LaTeX markup and return readable prose."""
    text = _COMMENT_RE.sub("", text)
    text = _ENV_EQUATION_RE.sub("", text)
    text = _ENV_ALIGN_RE.sub("", text)
    text = _DISPLAY_DOLLAR_RE.sub("", text)
    text = _INLINE_MATH_RE.sub("", text)
    text = _CITE_RE.sub("", text)
    text = _LABEL_RE.sub("", text)
    text = _REF_RE.sub("", text)
    text = _WHITESPACE_RE.sub(" ", text)
    text = _BLANK_LINES_RE.sub("\n\n", text)
    return text.strip()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def extract_sections(latex: str) -> list[tuple[str, str]]:
    """Split *latex* on \\section / \\subsection boundaries.

    Returns a list of ``(title, cleaned_text)`` tuples, one per section.
    Text between the document start and the first section command is
    discarded (preamble / abstract handling is out of scope).
    """
    # Find all section positions
    matches = list(_SECTION_RE.finditer(latex))
    if not matches:
        return []

    sections: list[tuple[str, str]] = []
    for idx, match in enumerate(matches):
        title = match.group(1).strip()
        body_start = match.end()
        body_end = matches[idx + 1].start() if idx + 1 < len(matches) else len(latex)
        body = latex[body_start:body_end]
        cleaned = _clean_latex(body)
        sections.append((title, cleaned))

    return sections


def chunk_text(text: str, size: int = 512, overlap: int = 64) -> list[str]:
    """Split *text* into overlapping word-based chunks.

    Parameters
    ----------
    text:
        Plain text to chunk (not raw LaTeX).
    size:
        Target number of words per chunk.
    overlap:
        Number of words to carry over into the next chunk.

    Each chunk is at most ``size + overlap`` words.  Boundaries are snapped
    to the nearest ``". "`` (sentence end) within ±20 words of the nominal
    boundary when possible.
    """
    words = text.split()
    if not words:
        return []

    snap_window = 20
    chunks: list[str] = []
    start = 0

    while start < len(words):
        end = min(start + size, len(words))

        # Try to snap *end* to a sentence boundary within ±snap_window words.
        if end < len(words):
            best = end
            # Build a small search window
            lo = max(start + 1, end - snap_window)
            hi = min(len(words), end + snap_window + 1)
            # Prefer the closest sentence-end ". " to *end*
            for offset in range(0, snap_window + 1):
                for candidate in (end - offset, end + offset):
                    if (
                        lo <= candidate < hi
                        and candidate > start
                        and words[candidate - 1].endswith(".")
                    ):
                        # Sentence boundary: word at (candidate-1) ends with ".".
                        best = candidate
                        break
                if best != end:
                    break
            end = best

        chunk_words = words[start:end]
        chunks.append(" ".join(chunk_words))

        if end >= len(words):
            break

        # Next chunk starts *overlap* words before *end*.
        start = max(start + 1, end - overlap)

    return chunks


def latex_to_text(latex: str) -> str:
    """Convert a LaTeX document to plain text.

    Extracts all sections, cleans each one, and joins them with a blank line.
    If no section commands are found, the whole document is cleaned and
    returned as a single block.
    """
    sections = extract_sections(latex)
    if sections:
        return "\n\n".join(body for _, body in sections)
    # Fallback: no section structure — clean the whole string.
    return _clean_latex(latex)