"""The output contract a `Parser` hands back: the document plus the signals
docling produces alongside it.

`result.document` is the DoclingDocument. Four more fields ride with it. The
fusion step reads only `pages` and `table_provenance`:

- `confidence` — docling's own confidence report, taken straight from
  `ConversionResult.confidence`. It is held in memory only. `save` does not
  write it and `load` leaves it `None`.
- `page_scores` — the per-page parse/layout/table/ocr scores copied out of that
  report. They are written to `<base>.confidence.json` and reloaded, but no code
  outside this module reads them.
- `pages` — the parsed-page cells (`text`, top-left box, `from_ocr`,
  `confidence`), one list per page. docling frees these after assembling the
  document unless `generate_parsed_pages=True`; capturing them is what makes the
  OCR-vs-native signal available downstream.
- `table_provenance` — derived once per table from the cells under its box: the
  fraction read by OCR, the mean OCR confidence, and a native/ocr/empty verdict.
  A table whose region carries no native text layer (verdict `ocr`, or `empty`
  when Camelot also finds nothing) is a table rendered as an image.

Parsed cells are stored top-left in PDF points. A table's box stays in the
DoclingDocument in docling's own coordinate origin. `_table_provenance` flips it
to top-left only while comparing it with the cells under it.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path

# Full docling appears in annotations only (`confidence`, `from_conversion`'s
# parameter) — never constructed here. Keeping it out of the runtime imports
# is what lets the CPU fuse job load artifacts without the GPU stack; the
# runtime needs docling-core alone.
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple

if TYPE_CHECKING:
    from docling.datamodel.base_models import ConfidenceReport
    from docling.datamodel.document import ConversionResult

from docling_core.types.doc.base import CoordOrigin
from docling_core.types.doc.document import DoclingDocument, TableItem
from docling_core.types.doc.page import BoundingRectangle
from pydantic import BaseModel, Field

Box = Tuple[float, float, float, float]


class ParsedCell(BaseModel):
    """One parsed-page text cell, box in top-left PDF points."""

    text: str
    box: Box = Field(description="(x0, y0, x1, y1), top-left origin, PDF points")
    from_ocr: bool = Field(description="True if docling read this cell off the page image")
    confidence: float = Field(default=0.0, description="docling's OCR confidence for the cell")


class PageParse(BaseModel):
    """The parsed-page cells for one page, with the page's point dimensions."""

    page_no: int
    width: float
    height: float
    cells: List[ParsedCell] = Field(default_factory=list)


class PageScore(BaseModel):
    """docling's per-page confidence scores (NaN where docling did not score)."""

    page_no: int
    parse_score: float
    layout_score: float
    table_score: float
    ocr_score: float


class TableProvenance(BaseModel):
    """OCR-vs-native verdict for one docling table, from the cells under its box."""

    self_ref: str = Field(description="The TableItem.self_ref this provenance describes")
    page_no: int
    cell_count: int = Field(description="Parsed-page cells whose center falls in the table box")
    from_ocr_fraction: float = Field(description="Fraction of those cells docling read by OCR")
    mean_ocr_confidence: float = Field(description="Mean confidence over the OCR-read cells")
    verdict: Literal["native", "ocr", "empty"] = Field(
        description=(
            "native: the table box has a real text layer. ocr: docling read it off "
            "the page image (a table rendered as an image). empty: no cells fall in "
            "the box at all."
        )
    )


@dataclass
class ParseResult:
    """Everything one parse produced: the document and docling's side signals.

    `document` is the canonical IR. The rest is docling's confidence report, the
    per-page scores, the parsed-page cells, and the per-table OCR/native verdict
    derived from those cells.
    """

    document: DoclingDocument
    confidence: Optional[ConfidenceReport] = None
    page_scores: List[PageScore] = field(default_factory=list)
    pages: List[PageParse] = field(default_factory=list)
    table_provenance: List[TableProvenance] = field(default_factory=list)

    def provenance_by_ref(self) -> Dict[str, TableProvenance]:
        """Per-table provenance keyed by `TableItem.self_ref`."""
        return {p.self_ref: p for p in self.table_provenance}

    def save(self, directory: Path, base: str) -> Dict[str, Path]:
        """Write the parse artifacts under `directory`, named from `base`.

        Three files form the document step's artifact contract — the seam the
        fusion step (and a later cloud job) reloads:

        - `<base>.docling.json` — the DoclingDocument.
        - `<base>.confidence.json` — per-page scores + per-table provenance.
        - `<base>.cells.json` — the parsed-page cells.

        Returns the written paths keyed by `document` / `confidence` / `cells`.
        """
        directory.mkdir(parents=True, exist_ok=True)
        document_path = directory / f"{base}.docling.json"
        confidence_path = directory / f"{base}.confidence.json"
        cells_path = directory / f"{base}.cells.json"

        self.document.save_as_json(document_path)
        confidence_path.write_text(
            json.dumps(
                {
                    "page_scores": [s.model_dump() for s in self.page_scores],
                    "table_provenance": [p.model_dump() for p in self.table_provenance],
                },
                indent=2,
                default=str,
            ),
            encoding="utf-8",
        )
        cells_path.write_text(
            json.dumps([p.model_dump() for p in self.pages], indent=2, default=str),
            encoding="utf-8",
        )
        return {"document": document_path, "confidence": confidence_path, "cells": cells_path}

    @classmethod
    def load(cls, directory: Path, base: str) -> "ParseResult":
        """Reconstruct a ParseResult from artifacts written by `save`.

        The `confidence` object is not restored. The document, page scores, page
        cells, and per-table provenance round-trip exactly. The fusion step reads
        the document, the page cells, and the per-table provenance.
        """
        document = DoclingDocument.load_from_json(directory / f"{base}.docling.json")
        confidence_blob = json.loads((directory / f"{base}.confidence.json").read_text())
        page_scores = [PageScore(**s) for s in confidence_blob.get("page_scores", [])]
        table_provenance = [TableProvenance(**p) for p in confidence_blob.get("table_provenance", [])]
        cells_blob = json.loads((directory / f"{base}.cells.json").read_text())
        pages = [PageParse(**p) for p in cells_blob]
        return cls(
            document=document,
            confidence=None,
            page_scores=page_scores,
            pages=pages,
            table_provenance=table_provenance,
        )

    @classmethod
    def from_conversion(cls, result: ConversionResult) -> "ParseResult":
        """Build a ParseResult from a docling ConversionResult.

        Reads the per-page scores and the parsed-page cells off the result, then
        derives the per-table provenance from the cells under each table box.
        Requires the converter to have run with `generate_parsed_pages=True`;
        without it `result.pages[i].cells` is empty and every table reads `empty`.
        """
        document = result.document

        page_scores: List[PageScore] = []
        if result.confidence is not None:
            for page_no, scores in sorted(result.confidence.pages.items()):
                page_scores.append(
                    PageScore(
                        page_no=page_no,
                        parse_score=scores.parse_score,
                        layout_score=scores.layout_score,
                        table_score=scores.table_score,
                        ocr_score=scores.ocr_score,
                    )
                )

        pages: List[PageParse] = []
        for page in result.pages:
            width = page.size.width if page.size else 0.0
            height = page.size.height if page.size else 0.0
            cells = [
                ParsedCell(
                    text=cell.text,
                    box=_cell_box_top_left(cell.rect),
                    from_ocr=cell.from_ocr,
                    confidence=cell.confidence or 0.0,
                )
                for cell in page.cells
            ]
            pages.append(PageParse(page_no=page.page_no, width=width, height=height, cells=cells))

        provenance = derive_table_provenance(document, pages)

        return cls(
            document=document,
            confidence=result.confidence,
            page_scores=page_scores,
            pages=pages,
            table_provenance=provenance,
        )


def derive_table_provenance(document: DoclingDocument, pages: List[PageParse]) -> List[TableProvenance]:
    """The OCR-versus-native verdict for every table in `document`.

    Purely geometric: each table's box is compared against the parsed-page cells
    it encloses. It runs once, in `ParseResult.from_conversion`, on the document
    docling just produced. Fusion looks the stored verdicts up by each table's
    parse-time `self_ref` through `provenance_by_ref`.
    """
    pages_by_no = {p.page_no: p for p in pages}
    return [_table_provenance(table, pages_by_no) for table in document.tables if table.prov]


def _cell_box_top_left(rect: BoundingRectangle) -> Box:
    """A parsed-cell rectangle (four corner points) -> a top-left (x0,y0,x1,y1) box."""
    xs = [rect.r_x0, rect.r_x1, rect.r_x2, rect.r_x3]
    ys = [rect.r_y0, rect.r_y1, rect.r_y2, rect.r_y3]
    return (min(xs), min(ys), max(xs), max(ys))


def _table_box_top_left(table: TableItem, page_height: float) -> Box:
    """A table's provenance bbox -> a top-left (x0,y0,x1,y1) box in PDF points.

    docling table boxes are bottom-left origin; flip y against the page height so
    they share the parsed cells' top-left frame.
    """
    bbox = table.prov[0].bbox
    if bbox.coord_origin == CoordOrigin.TOPLEFT:
        return (bbox.l, bbox.t, bbox.r, bbox.b)
    return (bbox.l, page_height - bbox.t, bbox.r, page_height - bbox.b)


def _center_in(box: Box, region: Box) -> bool:
    cx, cy = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2
    x0, y0, x1, y1 = (
        min(region[0], region[2]),
        min(region[1], region[3]),
        max(region[0], region[2]),
        max(region[1], region[3]),
    )
    return x0 <= cx <= x1 and y0 <= cy <= y1


def _table_provenance(table: TableItem, pages_by_no: Dict[int, PageParse]) -> TableProvenance:
    """Derive the OCR/native verdict for one table from the cells under its box."""
    page_no = table.prov[0].page_no
    page = pages_by_no.get(page_no)
    self_ref = table.self_ref

    if page is None:
        return TableProvenance(
            self_ref=self_ref,
            page_no=page_no,
            cell_count=0,
            from_ocr_fraction=0.0,
            mean_ocr_confidence=0.0,
            verdict="empty",
        )

    region = _table_box_top_left(table, page.height)
    under = [c for c in page.cells if _center_in(c.box, region)]
    if not under:
        return TableProvenance(
            self_ref=self_ref,
            page_no=page_no,
            cell_count=0,
            from_ocr_fraction=0.0,
            mean_ocr_confidence=0.0,
            verdict="empty",
        )

    ocr_cells = [c for c in under if c.from_ocr]
    from_ocr_fraction = len(ocr_cells) / len(under)
    mean_conf = (sum(c.confidence for c in ocr_cells) / len(ocr_cells)) if ocr_cells else 0.0
    verdict: Literal["native", "ocr", "empty"] = "ocr" if from_ocr_fraction > 0.5 else "native"
    return TableProvenance(
        self_ref=self_ref,
        page_no=page_no,
        cell_count=len(under),
        from_ocr_fraction=from_ocr_fraction,
        mean_ocr_confidence=mean_conf,
        verdict=verdict,
    )