"""OCR grounding well: page text read from a rendered image, with provenance.

A scanned page or an image-bound figure has no native text layer, so a value
extracted from it has nothing to ground against. The well fills that gap. It
runs RapidOCR over a rendered page image and keeps every recognized fragment
with its bounding box and confidence. For a render of an original PDF page the
box is stored in PDF points, so a fragment grounds a value from a scanned
region the same way a native text cell grounds a table cell.

The well is evidence, never a correction target. Fragments are stored exactly
as the engine read them; anything that interprets or repairs them happens in
the consumer, against the fragment ids, so the trace from a value back to the
pixels survives every downstream step.
"""

from __future__ import annotations

import subprocess
import tempfile
from pathlib import Path
from typing import Any

from pydantic import BaseModel

POINTS_PER_INCH = 72.0


class OcrFragment(BaseModel):
    """One recognized text fragment and where it sits on the page."""

    id: str
    page: int
    text: str
    #: Left, top, right, bottom. PDF points when the well was built from a PDF
    #: render; pixels when built from a standalone image (dpi is None then).
    bbox: tuple[float, float, float, float]
    confidence: float


class PageWell(BaseModel):
    """Every fragment read from one page image."""

    page: int
    image: str
    dpi: int | None
    width: float
    height: float
    fragments: list[OcrFragment]

    def fragments_in_box(
        self,
        bbox: tuple[float, float, float, float],
        min_overlap: float = 0.5,
    ) -> list[OcrFragment]:
        """Fragments whose area overlaps the region by at least min_overlap."""
        left, top, right, bottom = bbox
        hits = []
        for frag in self.fragments:
            fl, ft, fr, fb = frag.bbox
            inter_w = max(0.0, min(fr, right) - max(fl, left))
            inter_h = max(0.0, min(fb, bottom) - max(ft, top))
            area = max((fr - fl) * (fb - ft), 1e-6)
            if (inter_w * inter_h) / area >= min_overlap:
                hits.append(frag)
        return hits


class DocumentWell(BaseModel):
    """The wells for every page read from one source document."""

    source: str
    engine: str
    pages: list[PageWell]

    def page(self, page_no: int) -> PageWell:
        for well in self.pages:
            if well.page == page_no:
                return well
        raise KeyError(f"no well for page {page_no}")


def make_engine(device_id: int = 0) -> Any:
    """RapidOCR with the onnxruntime CUDA execution provider requested.

    docling propagates the CUDA device to RapidOCR's paddle and torch engines
    but not the default onnxruntime engine, so the CUDA execution provider is
    named explicitly here. This engine requests it on every host.
    TunedFinancialParser requests it only when the resolved accelerator device
    is CUDA.
    """
    from rapidocr import RapidOCR

    return RapidOCR(
        params={
            "EngineConfig.onnxruntime.use_cuda": True,
            "EngineConfig.onnxruntime.cuda_ep_cfg.device_id": device_id,
        }
    )


def read_image(
    image_path: Path,
    page_no: int,
    dpi: int | None = None,
    engine: Any = None,
) -> PageWell:
    """Read one page image into a well.

    When dpi is given the fragment boxes are converted from pixels to PDF
    points, so they land in the coordinate frame of the original page.
    """
    from PIL import Image

    if engine is None:
        engine = make_engine()
    with Image.open(image_path) as img:
        width_px, height_px = img.size
    result = engine(str(image_path))
    scale = POINTS_PER_INCH / dpi if dpi else 1.0

    fragments: list[OcrFragment] = []
    boxes = getattr(result, "boxes", None)
    txts = getattr(result, "txts", None) or []
    scores = getattr(result, "scores", None) or []
    if boxes is not None:
        for i, (box, text, score) in enumerate(zip(boxes, txts, scores, strict=False)):
            xs = [float(p[0]) for p in box]
            ys = [float(p[1]) for p in box]
            fragments.append(
                OcrFragment(
                    id=f"p{page_no}.f{i}",
                    page=page_no,
                    text=str(text),
                    bbox=(
                        min(xs) * scale,
                        min(ys) * scale,
                        max(xs) * scale,
                        max(ys) * scale,
                    ),
                    confidence=float(score),
                )
            )
    return PageWell(
        page=page_no,
        image=str(image_path),
        dpi=dpi,
        width=width_px * scale,
        height=height_px * scale,
        fragments=fragments,
    )


def build_well(
    pdf: Path,
    pages: list[int],
    dpi: int = 300,
    engine: Any = None,
) -> DocumentWell:
    """Render the named pages of a PDF and read each into a well."""
    if engine is None:
        engine = make_engine()
    page_wells: list[PageWell] = []
    with tempfile.TemporaryDirectory() as tmp:
        for page_no in pages:
            prefix = Path(tmp) / f"page{page_no:04d}"
            subprocess.run(
                [
                    "pdftoppm",
                    "-r",
                    str(dpi),
                    "-png",
                    "-f",
                    str(page_no),
                    "-l",
                    str(page_no),
                    str(pdf),
                    str(prefix),
                ],
                check=True,
                capture_output=True,
            )
            rendered = sorted(Path(tmp).glob(f"page{page_no:04d}*.png"))
            if not rendered:
                raise FileNotFoundError(f"pdftoppm produced no image for page {page_no}")
            page_wells.append(read_image(rendered[0], page_no, dpi=dpi, engine=engine))
    return DocumentWell(source=str(pdf), engine="rapidocr", pages=page_wells)