"""
PDF digital parser — uses pymupdf4llm for text + layout extraction.

Strategy: page-by-page processing to keep memory proportional to largest page,
not full document size.
"""

from __future__ import annotations

import importlib.util as _importlib_util
import logging
import re
from collections.abc import Iterator

from .base import BaseParser, FileRecord, RawElement

logger = logging.getLogger(__name__)

_HAS_SURYA = _importlib_util.find_spec("surya") is not None
_HAS_PDFPLUMBER = _importlib_util.find_spec("pdfplumber") is not None


_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$")
_CODE_FENCE_RE = re.compile(r"^```(\w*)$")
_TABLE_ROW_RE = re.compile(r"^\|.+\|$")
_TABLE_SEP_RE = re.compile(r"^\|[\s\-|:]+\|$")
_BULLET_RE = re.compile(r"^(?:[-*\u2022\u00b7]|\d+[.)]) +")


def _parse_table_markdown(lines: list[str]) -> dict:
    """Parse a GFM table from markdown lines into dual-format dict."""
    if len(lines) < 2:
        return {
            "markdown": "\n".join(lines),
            "structured": {"headers": [], "rows": []},
            "has_header_row": False,
            "row_count": len(lines),
            "header_row_count": 0,
            "col_count": 0,
        }

    def split_row(line: str) -> list[str]:
        return [cell.strip() for cell in line.strip().strip("|").split("|")]

    headers: list[str] = []
    rows: list[list[str]] = []
    has_header = False
    data_lines = lines

    # Detect header: row 0 + separator row 1
    if len(lines) >= 2 and _TABLE_SEP_RE.match(lines[1]):
        headers = split_row(lines[0])
        has_header = True
        data_lines = lines[2:]

    for line in data_lines:
        if _TABLE_ROW_RE.match(line) and not _TABLE_SEP_RE.match(line):
            rows.append(split_row(line))

    return {
        "markdown": "\n".join(lines),
        "structured": {"headers": headers, "rows": rows},
        "has_header_row": has_header,
        "row_count": len(rows),
        "header_row_count": 1 if has_header else 0,
        "col_count": len(headers) if headers else (len(rows[0]) if rows else 0),
    }


def _parse_markdown_page(
    md_text: str, page_num: int, extraction_method: str = "pymupdf4llm"
) -> Iterator[RawElement]:
    """Parse markdown text from a single page into RawElement stream."""
    lines = md_text.splitlines()
    i = 0
    while i < len(lines):
        line = lines[i]

        # Heading
        m = _HEADING_RE.match(line)
        if m:
            level = len(m.group(1))
            text = m.group(2).strip()
            yield RawElement(
                element_type="structural:section_header" if level > 1 else "structural:title",
                text=text,
                page=page_num,
                heading_level=min(level, 6),
                detection_tier=1,
                extra={"extraction_method": extraction_method, "confidence": 0.95},
            )
            i += 1
            continue

        # Code fence
        m_fence = _CODE_FENCE_RE.match(line)
        if m_fence:
            lang = m_fence.group(1)
            code_lines = []
            i += 1
            while i < len(lines) and not lines[i].startswith("```"):
                code_lines.append(lines[i])
                i += 1
            i += 1  # skip closing fence
            yield RawElement(
                element_type="code:block",
                text="\n".join(code_lines),
                page=page_num,
                extra={
                    "extraction_method": extraction_method,
                    "confidence": 0.98,
                    "language": lang or None,
                },
            )
            continue

        # Table — collect all table lines
        if _TABLE_ROW_RE.match(line):
            table_lines = []
            while i < len(lines) and (
                _TABLE_ROW_RE.match(lines[i]) or _TABLE_SEP_RE.match(lines[i])
            ):
                table_lines.append(lines[i])
                i += 1
            table_data = _parse_table_markdown(table_lines)
            yield RawElement(
                element_type="table:simple",
                text=table_data["markdown"],
                page=page_num,
                table_data=table_data,
                extra={"extraction_method": extraction_method, "confidence": 0.90},
            )
            continue

        # Paragraph / narrative text — detect bullet items
        stripped = line.strip()
        if stripped:
            # Detect bullet list items: lines starting with - / * / • / · or number.
            if _BULLET_RE.match(stripped):
                item_text = _BULLET_RE.sub("", stripped)
                yield RawElement(
                    element_type="list:item",
                    text=item_text,
                    page=page_num,
                    extra={"extraction_method": extraction_method, "confidence": 0.80},
                )
            else:
                yield RawElement(
                    element_type="text:narrative",
                    text=stripped,
                    page=page_num,
                    extra={"extraction_method": extraction_method, "confidence": 0.92},
                )

        i += 1


_layout_predictor_cache = None


def _get_layout_predictor():
    global _layout_predictor_cache
    if _layout_predictor_cache is None:
        from surya.layout import batch_layout_detection
        from surya.model.layout.model import load_model as load_layout_model
        from surya.model.layout.processor import load_processor as load_layout_proc

        _layout_predictor_cache = (batch_layout_detection, load_layout_model(), load_layout_proc())
    return _layout_predictor_cache


def _to_pdf_rect(bbox, img_size, page_rect):
    import pymupdf

    x0, y0, x1, y1 = bbox
    sx = page_rect.width / img_size[0]
    sy = page_rect.height / img_size[1]
    return pymupdf.Rect(x0 * sx, y0 * sy, x1 * sx, y1 * sy)


def _parse_page_with_surya(doc, page_num: int):
    from .pdf_ocr import _SURYA_TYPE_MAP, _page_to_pil

    batch_fn, model, processor = _get_layout_predictor()
    image = _page_to_pil(doc, page_num)
    layout_result = batch_fn([image], processor, model)
    page = doc[page_num]
    for box in sorted(layout_result[0].bboxes, key=lambda b: b.position):
        rect = _to_pdf_rect(box.bbox, image.size, page.rect)
        text = page.get_textbox(rect).strip()
        if not text:
            continue
        etype = _SURYA_TYPE_MAP.get(box.label, "text:narrative")
        yield RawElement(
            element_type=etype,
            text=text,
            page=page_num + 1,
            extra={"extraction_method": "surya:layout-digital", "confidence": 0.90},
        )


def _parse_page_with_pdfplumber(path, page_num: int, md_text: str):
    import pdfplumber

    with pdfplumber.open(str(path)) as pdf:
        pl_page = pdf.pages[page_num]
        words = pl_page.extract_words(x_tolerance=3, y_tolerance=3)
        page_width = pl_page.width
    if not words:
        yield from _parse_markdown_page(md_text, page_num + 1, "pdfplumber:heuristic")
        return
    xs = sorted({round(w["x0"]) for w in words})
    gutters = [xs[i] for i in range(1, len(xs)) if xs[i] - xs[i - 1] > page_width * 0.05]
    splits = [0] + gutters + [page_width + 1]
    cols: list[list] = [[] for _ in range(len(splits) - 1)]
    for w in words:
        for i in range(len(splits) - 1):
            if splits[i] <= w["x0"] < splits[i + 1]:
                cols[i].append(w)
                break
    text = "\n".join(
        " ".join(w["text"] for w in sorted(col, key=lambda w: (round(w["top"] / 5) * 5, w["x0"])))
        for col in cols
        if col
    )
    yield from _parse_markdown_page(text, page_num + 1, "pdfplumber:column-recovery")


def _extract_page_figures(
    doc: object,
    page_num: int,
    figures_dir: object,
) -> Iterator[RawElement]:
    """Extract embedded images from a PDF page and yield figure RawElements."""
    import fitz  # type: ignore[import-untyped]

    page = doc[page_num]  # type: ignore[index]
    try:
        image_list = page.get_images(full=True)
    except Exception:
        return
    for img_idx, img_info in enumerate(image_list):
        xref = img_info[0]
        try:
            pix = fitz.Pixmap(doc, xref)
            if pix.width * pix.height <= 100 * 100:
                pix.close()
                pix = None
                continue
            if pix.n > 4:  # CMYK -> RGB
                orig_pix = pix
                pix = fitz.Pixmap(fitz.csRGB, orig_pix)
                orig_pix.close()
                orig_pix = None
            fname = f"fig_p{page_num + 1}_{img_idx + 1}.png"
            asset_path = figures_dir / fname  # type: ignore[operator]
            pix.save(str(asset_path))
            yield RawElement(
                element_type="media:figure",
                text=f"[Figure: {fname}]",
                page=page_num + 1,
                figure_data={
                    "width_px": pix.width,
                    "height_px": pix.height,
                    "asset_path": str(asset_path),
                    "asset_format": "png",
                },
                extra={"extraction_method": "pymupdf:image", "confidence": 0.85},
            )
        except Exception as exc:
            logger.warning(
                "PDF figure extraction failed p%d img%d: %s", page_num + 1, img_idx + 1, exc
            )


class PDFDigitalParser(BaseParser):
    handles = [("application/pdf", "digital")]

    def parse(self, file_record: FileRecord) -> Iterator[RawElement]:
        yield from self.parse_with_figures(file_record, figures_dir=None)

    def parse_with_figures(
        self, file_record: FileRecord, figures_dir: object = None
    ) -> Iterator[RawElement]:
        try:
            import pymupdf  # type: ignore[import-untyped]
            import pymupdf4llm  # type: ignore[import-untyped]
        except ImportError as e:
            raise ImportError(f"pymupdf4llm required: pip install pymupdf4llm — {e}") from e

        doc = pymupdf.open(str(file_record.path))
        try:
            for page_num in range(doc.page_count):
                md_text = pymupdf4llm.to_markdown(doc, pages=[page_num])
                if _HAS_SURYA:
                    try:
                        yield from _parse_page_with_surya(doc, page_num)
                        if figures_dir is not None:
                            yield from _extract_page_figures(doc, page_num, figures_dir)
                        continue
                    except Exception as _exc:
                        logger.warning(
                            "PDF page %d parse failed (%s tier): %s", page_num, "surya", _exc
                        )
                if _HAS_PDFPLUMBER:
                    try:
                        if not md_text or not md_text.strip():
                            if figures_dir is not None:
                                yield from _extract_page_figures(doc, page_num, figures_dir)
                            continue
                        yield from _parse_page_with_pdfplumber(file_record.path, page_num, md_text)
                        if figures_dir is not None:
                            yield from _extract_page_figures(doc, page_num, figures_dir)
                        continue
                    except Exception as _exc:
                        logger.warning(
                            "PDF page %d parse failed (%s tier): %s", page_num, "pdfplumber", _exc
                        )
                if not md_text or not md_text.strip():
                    if figures_dir is not None:
                        yield from _extract_page_figures(doc, page_num, figures_dir)
                    continue
                yield from _parse_markdown_page(md_text, page_num + 1, "pymupdf4llm")
                if figures_dir is not None:
                    yield from _extract_page_figures(doc, page_num, figures_dir)
        finally:
            doc.close()
