"""
Scanned PDF OCR parser — uses surya for layout + OCR.

Strategy:
1. Convert each PDF page to image (pymupdf at 150 DPI)
2. Run surya layout detection to find text regions, figures, tables
3. Run surya OCR on each text region
4. Yield RawElements based on surya layout block types:
   - SectionHeader → structural:section_header
   - Title → structural:title
   - Text → text:narrative
   - Table → table:simple (OCR text only, no structure parsing)
   - Figure → media:figure (caption from nearby Text block)
   - Formula → scientific:equation_display
   - Code → code:block

Falls back to raw pymupdf text extraction if surya is not installed.
"""

from __future__ import annotations

from collections.abc import Iterator

from .base import BaseParser, FileRecord, RawElement
from .pdf_digital import _parse_markdown_page

# Surya layout block type → element type mapping
_SURYA_TYPE_MAP: dict[str, str] = {
    "Title": "structural:title",
    "SectionHeader": "structural:section_header",
    "Text": "text:narrative",
    "Caption": "text:caption",
    "Footnote": "text:footnote",
    "Formula": "scientific:equation_display",
    "Code": "code:block",
    "Table": "table:simple",
    "Figure": "media:figure",
    "ListItem": "list:item",
    "PageHeader": "structural:page_header",
    "PageFooter": "structural:page_footer",
}

# Resolution for page-to-image conversion
_DPI = 150


def _page_to_pil(doc, page_num: int):  # type: ignore[no-untyped-def]
    """Render a pymupdf page to a PIL Image at _DPI."""
    import io

    from PIL import Image  # type: ignore[import-untyped]

    page = doc[page_num]
    mat = page.get_pixmap(dpi=_DPI)
    img_bytes = mat.tobytes("png")
    return Image.open(io.BytesIO(img_bytes)).convert("RGB")


def _ocr_page(image, page_num: int) -> Iterator[RawElement]:  # type: ignore[no-untyped-def]
    """Run surya layout + OCR on a single PIL image, yield RawElements."""
    try:
        from surya.model.detection.model import (
            load_model as load_det_model,  # type: ignore[import-untyped]
        )
        from surya.model.detection.processor import (
            load_processor as load_det_proc,  # type: ignore[import-untyped]
        )
        from surya.model.recognition.model import (
            load_model as load_rec_model,  # type: ignore[import-untyped]
        )
        from surya.model.recognition.processor import (
            load_processor as load_rec_proc,  # type: ignore[import-untyped]
        )
        from surya.ocr import run_ocr  # type: ignore[import-untyped]
    except ImportError as e:
        raise ImportError(f"surya-ocr required: pip install surya-ocr — {e}") from e

    # Load models (cached after first call via surya's internal caching)
    det_model = load_det_model()
    det_proc = load_det_proc()
    rec_model = load_rec_model()
    rec_proc = load_rec_proc()

    # Run OCR — surya returns line-level results
    results = run_ocr([image], [["en"]], det_model, det_proc, rec_model, rec_proc)
    page_result = results[0] if results else None
    if not page_result:
        return

    # Aggregate OCR lines into block text
    for line in page_result.text_lines:
        text = line.text.strip()
        if not text:
            continue
        yield RawElement(
            element_type="text:narrative",
            text=text,
            page=page_num,
            extra={"extraction_method": "surya:ocr", "confidence": float(line.confidence)},
        )


def _ocr_page_with_layout(image, page_num: int) -> Iterator[RawElement]:  # type: ignore[no-untyped-def]
    """Run surya layout detection + OCR, respecting block types."""
    try:
        from surya.layout import batch_layout_detection  # type: ignore[import-untyped]
        from surya.model.detection.model import (
            load_model as load_det_model,  # type: ignore[import-untyped]
        )
        from surya.model.detection.processor import (
            load_processor as load_det_proc,  # type: ignore[import-untyped]
        )
        from surya.model.layout.model import (
            load_model as load_layout_model,  # type: ignore[import-untyped]
        )
        from surya.model.layout.processor import (
            load_processor as load_layout_proc,  # type: ignore[import-untyped]
        )
        from surya.model.recognition.model import (
            load_model as load_rec_model,  # type: ignore[import-untyped]
        )
        from surya.model.recognition.processor import (
            load_processor as load_rec_proc,  # type: ignore[import-untyped]
        )
        from surya.ocr import run_ocr  # type: ignore[import-untyped]
    except ImportError:
        # Surya layout not available — fall back to line-only OCR
        yield from _ocr_page(image, page_num)
        return

    layout_model = load_layout_model()
    layout_proc = load_layout_proc()
    det_model = load_det_model()
    det_proc = load_det_proc()
    rec_model = load_rec_model()
    rec_proc = load_rec_proc()

    # Layout detection
    layout_results = batch_layout_detection([image], layout_model, layout_proc)
    layout_page = layout_results[0] if layout_results else None

    # OCR on full page
    ocr_results = run_ocr([image], [["en"]], det_model, det_proc, rec_model, rec_proc)
    ocr_page = ocr_results[0] if ocr_results else None

    if not layout_page or not ocr_page:
        return

    # Build a list of (bbox, text) from OCR lines
    def bbox_intersects(b1, b2) -> bool:  # type: ignore[no-untyped-def]
        """Check if two bboxes overlap (each is [x0,y0,x1,y1])."""
        return not (b1[2] < b2[0] or b1[0] > b2[2] or b1[3] < b2[1] or b1[1] > b2[3])

    ocr_lines = [
        (line.bbox, line.text.strip(), line.confidence)
        for line in ocr_page.text_lines
        if line.text.strip()
    ]

    # Assign OCR lines to layout blocks
    for block in sorted(layout_page.bboxes, key=lambda b: b.position):
        block_type = block.label
        element_type = _SURYA_TYPE_MAP.get(block_type, "text:narrative")
        block_bbox = block.bbox

        # Collect OCR lines that fall within this block
        block_lines = [
            (txt, conf) for bbox, txt, conf in ocr_lines if bbox_intersects(bbox, block_bbox)
        ]
        if not block_lines:
            continue

        text = " ".join(t for t, _ in block_lines).strip()
        confidence = sum(c for _, c in block_lines) / len(block_lines)

        if not text:
            continue

        extra: dict = {
            "extraction_method": "surya:layout+ocr",
            "confidence": round(confidence, 4),
            "layout_block_type": block_type,
        }

        if element_type == "structural:title":
            yield RawElement(
                element_type=element_type,
                text=text,
                page=page_num,
                heading_level=1,
                detection_tier=2,  # font heuristic via layout model
                extra=extra,
            )
        elif element_type == "structural:section_header":
            yield RawElement(
                element_type=element_type,
                text=text,
                page=page_num,
                heading_level=2,
                detection_tier=2,
                extra=extra,
            )
        elif element_type in ("table:simple",):
            yield RawElement(
                element_type=element_type,
                text=text,
                page=page_num,
                table_data={
                    "markdown": text,
                    "structured": {"headers": [], "rows": []},
                    "has_header_row": False,
                    "row_count": 0,
                    "header_row_count": 0,
                    "col_count": 0,
                    "note": "OCR-only; no structure parsed",
                },
                extra=extra,
            )
        elif element_type == "media:figure":
            yield RawElement(
                element_type=element_type,
                text=text,
                page=page_num,
                figure_data={
                    "caption": text,
                    "image_ref": None,
                    "image_sha256": None,
                    "ocr_text": text,
                },
                extra=extra,
            )
        elif element_type == "scientific:equation_display":
            yield RawElement(
                element_type=element_type,
                text=text,
                page=page_num,
                equation_data={"latex": None, "label": None, "ocr_text": text},
                extra=extra,
            )
        else:
            yield RawElement(
                element_type=element_type,
                text=text,
                page=page_num,
                extra=extra,
            )


def _pymupdf_text_fallback(path) -> Iterator[RawElement]:  # type: ignore[no-untyped-def]
    """Last-resort fallback: raw pymupdf text extraction page by page."""
    try:
        import pymupdf  # type: ignore[import-untyped]
        import pymupdf4llm  # type: ignore[import-untyped]
    except ImportError as e:
        raise ImportError(f"pymupdf required: pip install pymupdf — {e}") from e

    doc = pymupdf.open(str(path))
    try:
        for page_num in range(doc.page_count):
            md_text = pymupdf4llm.to_markdown(doc, pages=[page_num])
            if md_text and md_text.strip():
                yield from _parse_markdown_page(md_text, page_num + 1, "pymupdf:scanned-fallback")
    finally:
        doc.close()


class PDFOCRParser(BaseParser):
    handles = [("application/pdf", "scanned")]

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

        # Check if surya is available
        try:
            import surya  # type: ignore[import-untyped]  # noqa: F401

            has_surya = True
        except ImportError:
            has_surya = False

        if not has_surya:
            yield from _pymupdf_text_fallback(file_record.path)
            return

        doc = pymupdf.open(str(file_record.path))
        try:
            for page_num in range(doc.page_count):
                try:
                    image = _page_to_pil(doc, page_num)
                    yield from _ocr_page_with_layout(image, page_num + 1)
                except Exception:
                    # Page-level failure: emit what we can from pymupdf text
                    import pymupdf4llm  # type: ignore[import-untyped]

                    md_text = pymupdf4llm.to_markdown(doc, pages=[page_num])
                    if md_text and md_text.strip():
                        yield from _parse_markdown_page(
                            md_text, page_num + 1, "pymupdf:page-ocr-fallback"
                        )
        finally:
            doc.close()
