"""
LaTeX source parser — walks TeX source and maps commands to RawElements.

Handles:
- \\title, \\author, \\date → meta:document_title / meta:author / meta:date
- \\section, \\subsection, \\subsubsection → structural:section_header/title
- Paragraph blocks (separated by blank lines) → text:narrative
- \\begin{abstract}...\\end{abstract} → text:abstract
- \\begin{equation}...\\end{equation}, $...$, $$...$$ → scientific:equation_display / scientific:equation_inline
- \\begin{table}...\\end{table} → table:simple (caption + raw content)
- \\begin{figure}...\\end{figure} → media:figure (caption)
- \\begin{verbatim}...\\end{verbatim}, \\lstlisting → code:block
- \\cite{...} annotations on narrative elements
- \\begin{itemize/enumerate}...\\item... → list:item / list:item_ordered
"""

from __future__ import annotations

import re
from collections.abc import Iterator

from .base import BaseParser, FileRecord, RawElement

# --- Patterns ---
_SECTION_RE = re.compile(
    r"\\(part|chapter|section|subsection|subsubsection|paragraph|subparagraph)"
    r"\*?\s*\{([^}]*)\}",
    re.DOTALL,
)
_TITLE_RE = re.compile(r"\\title\s*\{([^}]*)\}", re.DOTALL)
_AUTHOR_RE = re.compile(r"\\author\s*\{([^}]*)\}", re.DOTALL)
_LABEL_RE = re.compile(r"\\label\s*\{([^}]*)\}")
_CITE_RE = re.compile(r"\\cite[tp]?\s*(?:\[[^\]]*\])?\s*\{([^}]*)\}")
_CMD_RE = re.compile(r"\\[a-zA-Z]+\*?\s*(?:\[[^\]]*\])?\s*(?:\{[^}]*\})*")
_DISPLAY_MATH_BLOCK = re.compile(r"\$\$(.+?)\$\$", re.DOTALL)
_DISPLAY_MATH_ENV = re.compile(
    r"\\begin\{(equation\*?|align\*?|eqnarray\*?|multline\*?|gather\*?)\}(.+?)\\end\{\1\}",
    re.DOTALL,
)
_INLINE_MATH = re.compile(r"(?<!\$)\$(?!\$)(.+?)(?<!\$)\$(?!\$)", re.DOTALL)
_VERBATIM_ENV = re.compile(
    r"\\begin\{(verbatim|lstlisting|minted|Verbatim)\}(.*?)\\end\{\1\}",
    re.DOTALL,
)
_ENV_RE = re.compile(r"\\begin\{(\w+)\}(.*?)\\end\{\1\}", re.DOTALL)
_ITEM_RE = re.compile(r"\\item\s*(?:\[[^\]]*\])?\s*", re.DOTALL)

# Section command → heading level
_SECTION_LEVELS: dict[str, int] = {
    "part": 1,
    "chapter": 1,
    "section": 2,
    "subsection": 3,
    "subsubsection": 4,
    "paragraph": 5,
    "subparagraph": 6,
}


def _strip_latex(text: str) -> str:
    """Remove most LaTeX commands, leaving plain text."""
    # Remove comments
    text = re.sub(r"%.*", "", text)
    # Unwrap inline math $...$ by keeping the inner content visible
    # (prevents $\mathcal{C}$ → "$$" when the command is stripped)
    text = re.sub(r"(?<!\$)\$(?!\$)(.+?)(?<!\$)\$(?!\$)", lambda m: m.group(1), text)
    # Remove display math $$ blocks (already handled upstream as equation_display)
    text = re.sub(r"\$\$.*?\$\$", "", text, flags=re.DOTALL)
    # Remove known formatting commands but keep their argument
    text = re.sub(
        r"\\(?:textbf|textit|emph|texttt|textsc|underline|textrm)\{([^}]*)\}", r"\1", text
    )
    # Remove cite/ref/label
    text = re.sub(r"\\(?:cite|ref|label|eqref|autoref)\s*(?:\[[^\]]*\])?\s*\{[^}]*\}", "", text)
    # Remove all other commands
    text = re.sub(r"\\[a-zA-Z]+\*?(?:\[[^\]]*\])?(?:\{[^}]*\})*", "", text)
    text = re.sub(r"[{}]", "", text)
    return re.sub(r"\s+", " ", text).strip()


def _extract_caption(env_body: str) -> str:
    """Extract \\caption{...} text from an environment body."""
    m = re.search(r"\\caption\s*(?:\[[^\]]*\])?\s*\{([^}]*)\}", env_body, re.DOTALL)
    if m:
        return _strip_latex(m.group(1)).strip()
    return ""


def _extract_citations(text: str) -> list[str]:
    """Return list of citation keys found in text."""
    keys = []
    for m in _CITE_RE.finditer(text):
        keys.extend(k.strip() for k in m.group(1).split(","))
    return keys


def parse_latex(source: str) -> Iterator[RawElement]:  # noqa: C901
    """Walk LaTeX source string and yield RawElements in document order."""

    # Extract preamble metadata (before \\begin{document})
    doc_start = source.find(r"\begin{document}")
    preamble = source[:doc_start] if doc_start != -1 else ""
    body = source[doc_start + len(r"\begin{document}") :] if doc_start != -1 else source
    # Strip \\end{document}
    body = re.sub(r"\\end\{document\}.*", "", body, flags=re.DOTALL)

    # Title from preamble
    m = _TITLE_RE.search(preamble)
    if m:
        title = _strip_latex(m.group(1))
        if title:
            yield RawElement(
                element_type="structural:title",
                text=title,
                heading_level=1,
                detection_tier=1,
                extra={"extraction_method": "latex:preamble", "confidence": 0.99},
            )

    # Author from preamble
    m = _AUTHOR_RE.search(preamble)
    if m:
        # Replace \and with " and " before stripping so author names stay separated
        raw_author = re.sub(r"\\and\b", " and ", m.group(1))
        author = _strip_latex(raw_author)
        if author:
            yield RawElement(
                element_type="meta:author",
                text=author,
                extra={"extraction_method": "latex:preamble", "confidence": 0.99},
            )

    # We'll do a single-pass walk over the body, consuming environments greedily
    # in document order. We track position to interleave paragraph text.
    pos = 0
    # Collect all "interesting" spans in order
    spans: list[tuple[int, int, str, re.Match]] = []  # (start, end, kind, match)

    for m in _SECTION_RE.finditer(body):
        spans.append((m.start(), m.end(), "section", m))
    for m in _DISPLAY_MATH_BLOCK.finditer(body):
        spans.append((m.start(), m.end(), "display_math_block", m))
    for m in _DISPLAY_MATH_ENV.finditer(body):
        spans.append((m.start(), m.end(), "display_math_env", m))
    for m in _VERBATIM_ENV.finditer(body):
        spans.append((m.start(), m.end(), "verbatim", m))
    for m in _ENV_RE.finditer(body):
        spans.append((m.start(), m.end(), "env", m))

    # Sort by position; for overlaps keep the one that starts earliest (then longest)
    spans.sort(key=lambda x: (x[0], -(x[1] - x[0])))

    # Remove overlapping spans (keep earliest start, longest)
    clean_spans: list[tuple[int, int, str, re.Match]] = []
    last_end = 0
    for span in spans:
        if span[0] >= last_end:
            clean_spans.append(span)
            last_end = span[1]

    def _emit_text_block(text: str) -> Iterator[RawElement]:
        """Emit paragraph blocks from raw LaTeX text."""
        # Split on blank lines
        for block in re.split(r"\n\s*\n", text):
            plain = _strip_latex(block)
            if len(plain) < 10:
                continue
            citations = _extract_citations(block)
            extra: dict = {"extraction_method": "latex:paragraph", "confidence": 0.88}
            if citations:
                extra["inline_citations"] = citations
            yield RawElement(element_type="text:narrative", text=plain, extra=extra)

    for start, end, kind, m in clean_spans:
        # Emit any text before this span
        gap = body[pos:start]
        if gap.strip():
            yield from _emit_text_block(gap)
        pos = end

        if kind == "section":
            cmd = m.group(1)
            text = _strip_latex(m.group(2))
            level = _SECTION_LEVELS.get(cmd, 2)
            if text:
                yield RawElement(
                    element_type="structural:title" if level == 1 else "structural:section_header",
                    text=text,
                    heading_level=level,
                    detection_tier=1,
                    extra={"extraction_method": "latex:section", "confidence": 0.99},
                )

        elif kind == "display_math_block":
            latex = m.group(1).strip()
            if latex:
                yield RawElement(
                    element_type="scientific:equation_display",
                    text=latex,
                    equation_data={"latex": latex, "label": None},
                    extra={"extraction_method": "latex:display-math", "confidence": 0.97},
                )

        elif kind == "display_math_env":
            env_name = m.group(1)
            latex = m.group(2).strip()
            # Extract label if present
            label_m = _LABEL_RE.search(latex)
            label = label_m.group(1) if label_m else None
            latex_clean = _LABEL_RE.sub("", latex).strip()
            if latex_clean:
                yield RawElement(
                    element_type="scientific:equation_display",
                    text=latex_clean,
                    equation_data={"latex": latex_clean, "label": label, "env": env_name},
                    extra={"extraction_method": "latex:equation-env", "confidence": 0.97},
                )

        elif kind == "verbatim":
            code = m.group(2)
            env_name = m.group(1)
            # For lstlisting/minted, try to detect language from options
            lang = None
            opts = re.search(r"\\begin\{" + re.escape(env_name) + r"\}\[([^\]]*)\]", m.group(0))
            if opts:
                lang_m = re.search(r"language\s*=\s*(\w+)", opts.group(1))
                lang = lang_m.group(1) if lang_m else None
            if code.strip():
                yield RawElement(
                    element_type="code:block",
                    text=code.strip(),
                    extra={
                        "extraction_method": "latex:verbatim",
                        "confidence": 0.99,
                        "language": lang,
                    },
                )

        elif kind == "env":
            env_name = m.group(1)
            env_body = m.group(2)

            if env_name == "abstract":
                text = _strip_latex(env_body)
                if text:
                    yield RawElement(
                        element_type="text:abstract",
                        text=text,
                        extra={"extraction_method": "latex:abstract", "confidence": 0.99},
                    )

            elif env_name in ("table", "table*"):
                caption = _extract_caption(env_body)
                # Try to find tabular environment
                tabular_m = re.search(
                    r"\\begin\{tabular[x*]?\}\{[^}]*\}(.*?)\\end\{tabular[x*]?\}",
                    env_body,
                    re.DOTALL,
                )
                if tabular_m:
                    rows = _parse_tabular(tabular_m.group(1))
                    headers = rows.pop(0) if rows else []
                    md_lines = []
                    if headers:
                        md_lines.append("| " + " | ".join(headers) + " |")
                        md_lines.append("|" + "|".join(" --- " for _ in headers) + "|")
                    for row in rows:
                        padded = list(row) + [""] * max(0, len(headers) - len(row))
                        md_lines.append("| " + " | ".join(padded) + " |")
                    markdown = "\n".join(md_lines)
                    yield RawElement(
                        element_type="table:simple",
                        text=markdown or caption,
                        table_data={
                            "markdown": markdown,
                            "structured": {"headers": headers, "rows": rows},
                            "has_header_row": bool(headers),
                            "row_count": len(rows),
                            "header_row_count": 1 if headers else 0,
                            "col_count": len(headers) if headers else (len(rows[0]) if rows else 0),
                            "caption": caption or None,
                        },
                        extra={"extraction_method": "latex:table", "confidence": 0.92},
                    )
                elif caption:
                    yield RawElement(
                        element_type="table:simple",
                        text=caption,
                        table_data={
                            "markdown": "",
                            "structured": {"headers": [], "rows": []},
                            "has_header_row": False,
                            "row_count": 0,
                            "header_row_count": 0,
                            "col_count": 0,
                            "caption": caption,
                        },
                        extra={"extraction_method": "latex:table", "confidence": 0.80},
                    )

            elif env_name in ("figure", "figure*"):
                caption = _extract_caption(env_body)
                if caption:
                    yield RawElement(
                        element_type="media:figure",
                        text=caption,
                        figure_data={
                            "caption": caption,
                            "image_ref": None,
                            "image_sha256": None,
                            "ocr_text": None,
                        },
                        extra={"extraction_method": "latex:figure", "confidence": 0.95},
                    )

            elif env_name in ("itemize", "description"):
                for item_text in _ITEM_RE.split(env_body)[1:]:  # skip text before first \item
                    plain = _strip_latex(item_text)
                    if plain:
                        yield RawElement(
                            element_type="list:item",
                            text=plain,
                            extra={"extraction_method": "latex:itemize", "confidence": 0.90},
                        )

            elif env_name == "enumerate":
                for idx, item_text in enumerate(_ITEM_RE.split(env_body)[1:]):
                    plain = _strip_latex(item_text)
                    if plain:
                        yield RawElement(
                            element_type="list:item_ordered",
                            text=plain,
                            list_data={
                                "list_index": idx,
                                "list_depth": 0,
                                "first_item_sequence_index": 0,
                            },
                            extra={"extraction_method": "latex:enumerate", "confidence": 0.90},
                        )

            elif env_name == "thebibliography":
                for i, bibitem in enumerate(
                    re.split(r"\\bibitem\s*(?:\[[^\]]*\])?\s*\{[^}]*\}", env_body)[1:]
                ):
                    ref_text = _strip_latex(bibitem)
                    if ref_text:
                        yield RawElement(
                            element_type="scientific:reference_entry",
                            text=ref_text,
                            reference_data={"index": i + 1, "raw": ref_text},
                            extra={"extraction_method": "latex:bibliography", "confidence": 0.92},
                        )

            # Other environments: recurse into body as text
            else:
                text = _strip_latex(env_body)
                if len(text) >= 10:
                    yield RawElement(
                        element_type="text:narrative",
                        text=text,
                        extra={"extraction_method": "latex:env-fallback", "confidence": 0.75},
                    )

    # Emit any trailing text after last span
    tail = body[pos:]
    if tail.strip():
        yield from _emit_text_block(tail)


def _parse_tabular(tabular_body: str) -> list[list[str]]:
    """Parse LaTeX tabular body into list of rows (list of cell strings)."""
    rows: list[list[str]] = []
    # Remove \\hline and \\cline
    tabular_body = re.sub(r"\\hline|\\cline\{[^}]*\}", "", tabular_body)
    for row_str in re.split(r"\\\\", tabular_body):
        row_str = row_str.strip()
        if not row_str:
            continue
        cells = [_strip_latex(c) for c in row_str.split("&")]
        cells = [c for c in cells if c]
        if cells:
            rows.append(cells)
    return rows


class LaTeXParser(BaseParser):
    handles = [("text/x-tex", None), ("application/x-tex", None)]

    def parse(self, file_record: FileRecord) -> Iterator[RawElement]:
        try:
            from charset_normalizer import from_path

            result = from_path(str(file_record.path)).best()
            source = (
                str(result)
                if result
                else file_record.path.read_text(encoding="utf-8", errors="replace")
            )
        except Exception:
            source = file_record.path.read_text(encoding="utf-8", errors="replace")

        yield from parse_latex(source)
