"""Which pages get scanned, read off the parse.

ADE is the right reader for three things a financial document prints on a page,
and nomination is how we aim it at them.

It chooses pages, not regions. ADE locates and bounds whatever it finds on the
page it is handed, so there is nothing here to bound and no second judgement to
make. The only decision is where to send it.

The parse already reports the two picture signals, and the table signal is a
count of words in the source's text layer, so nomination makes no model call:

- A picture the parse classified as a chart. Plotted values are the whole point
  of a chart, and no component in the pipeline reads them.
- A picture the parse classified as something else with content in it — a
  photograph, a map, a diagram, a calendar. Also read by nothing today.
- A table with no text layer under it. Its cells hold text, so something read
  them, and the only thing that could have is the parse reading the page image.
  Such a table never reached the Set-of-Mark and Camelot engine either, because
  that engine extracts from a text layer and there was none, so what stands in
  the document is the parse's own reading with nothing having checked it since.
  `table_verdicts` measures this against the source document rather than against
  a record written earlier, so it holds however the tables have been renumbered
  since and needs no artifact beyond the two the workflow already takes.

Logos and icons are left where they are. A wordmark is printed on every page of a
deck as identity rather than content, and there is nothing inside it to read.

Every nominated page is scanned; nothing filters the list after it is built.
Nomination is a cost optimization, not a guarantee: a page whose content the
parse never classified — an infographic drawn entirely in the text layer — has
no signal here and is invisible to it. The orchestrator's `pages="all"` mode
exists for exactly those documents, scanning every page so nothing depends on
the parse having seen what matters.
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Mapping, Optional

from docling_core.types.doc.base import CoordOrigin
from docling_core.types.doc.document import DoclingDocument, PictureItem, TableItem
from loguru import logger
from pydantic import BaseModel, Field

from quber.agents.completeness import page_words

#: Picture classes the parse assigns to page furniture — the same mark repeated
#: on every page rather than document content. A scan reads them as readily as it
#: reads a chart, and there is nothing inside them to read.
FURNITURE_CLASSES = frozenset({"logo", "icon"})

#: Table verdicts that mean no text layer stood under the table's box, so what
#: the document holds is the parse's own reading of the page image. `empty` is
#: accepted for a caller whose verdict map marks a table with no cells that way.
#: `table_verdicts` never writes it: it emits only `native` or `ocr`, and a table
#: with no filled cells comes out `native`.
UNREAD_VERDICTS = frozenset({"ocr", "empty"})

#: Text-layer words under a table's box, as a fraction of the cells the table
#: holds, below which the table was not read from a text layer. A table typeset
#: as text carries at least one word per filled cell and usually several; a table
#: printed as an image carries none at all, so the margin between the two is
#: wide and the exact fraction does not decide any real case.
TEXT_LAYER_RATIO = 0.5


class NominatedPage(BaseModel):
    """One page the parse says holds something nothing has read."""

    page: int = Field(description="1-based source page")
    picture_refs: List[str] = Field(
        default_factory=list, description="The pictures on the page that nominated it, in document order"
    )
    picture_classes: List[str] = Field(
        default_factory=list, description="The class the parse assigned each of them"
    )
    table_refs: List[str] = Field(
        default_factory=list,
        description="Tables on the page the parse read off the page image, in document order",
    )


def nominate_pages(
    document: DoclingDocument,
    table_verdicts: Optional[Mapping[str, str]] = None,
) -> List[NominatedPage]:
    """Pages of `document` holding a picture or a table nothing has read.

    `table_verdicts` maps a table's `self_ref` to a verdict on whether its text
    stands in the page's text layer. The orchestrator passes the map this
    module's `table_verdicts` measures against the source document. A table
    nominates its page when its verdict is in `UNREAD_VERDICTS`. Omitted, no
    table nominates a page and only the pictures do.
    """
    verdicts = table_verdicts or {}
    by_page: dict[int, NominatedPage] = {}

    def page_entry(page: int) -> NominatedPage:
        return by_page.setdefault(page, NominatedPage(page=page))

    for picture in document.pictures:
        if not picture.prov:
            continue
        classes = picture_classes(picture)
        if classes and all(c in FURNITURE_CLASSES for c in classes):
            continue
        nominated = page_entry(picture.prov[0].page_no)
        nominated.picture_refs.append(picture.self_ref)
        nominated.picture_classes.extend(classes)

    for table in document.tables:
        if not table.prov:
            continue
        if verdicts.get(table.self_ref) not in UNREAD_VERDICTS:
            continue
        page_entry(table.prov[0].page_no).table_refs.append(table.self_ref)

    return [by_page[p] for p in sorted(by_page)]


def table_verdicts(document: DoclingDocument, source: Path) -> Dict[str, str]:
    """Whether each table's text stands in the page's text layer, keyed by reference.

    A table is `native` when the page prints text under its box, and `ocr` when
    it holds cells the page has no text for — the parse read those off the page
    image. The comparison is a count of words under the box against the cells the
    table holds, and the two cases are far apart: a table typeset as text carries
    at least a word per filled cell, one printed as an image carries none.

    A page whose text cannot be read at all gives its tables no entry in the
    map. `nominate_pages` treats a missing entry as read, so a failure here
    never sends a page to be scanned on a signal nobody measured.
    """
    by_page: Dict[int, List[TableItem]] = {}
    for table in document.tables:
        if table.prov:
            by_page.setdefault(table.prov[0].page_no, []).append(table)

    verdicts: Dict[str, str] = {}
    for page, tables in sorted(by_page.items()):
        try:
            _width, height, words = page_words(source, page)
        except Exception as exc:
            logger.warning(
                "Nomination: page {} text layer unreadable ({}: {}); its {} table(s) are taken as "
                "typeset text and none is nominated",
                page,
                type(exc).__name__,
                exc,
                len(tables),
            )
            continue
        for table in tables:
            filled = sum(1 for c in table.data.table_cells if (c.text or "").strip())
            under = sum(1 for w in words if _center_in(w, _top_left_box(table, height)))
            verdicts[table.self_ref] = "native" if under >= TEXT_LAYER_RATIO * filled else "ocr"
    return verdicts


def _top_left_box(table: TableItem, page_height: float) -> tuple[float, float, float, float]:
    """A table's provenance box as (left, top, right, bottom) in top-left points."""
    bbox = table.prov[0].bbox
    if bbox.coord_origin == CoordOrigin.TOPLEFT:
        return (bbox.l, min(bbox.t, bbox.b), bbox.r, max(bbox.t, bbox.b))
    return (bbox.l, page_height - max(bbox.t, bbox.b), bbox.r, page_height - min(bbox.t, bbox.b))


def _center_in(word: tuple[float, float, float, float, str], box: tuple[float, float, float, float]) -> bool:
    """True when a word's centre falls inside a top-left box."""
    left, top, right, bottom = box
    return left <= (word[0] + word[2]) / 2 <= right and top <= (word[1] + word[3]) / 2 <= bottom


def picture_classes(picture: PictureItem) -> List[str]:
    """The parse's top predicted class for a picture (empty if unclassified).

    Predictions are ordered by descending confidence, so the first entry is the
    call. Reads the current `meta.classification` field, falling back to the
    deprecated `annotations` list for output written before that move.
    """
    meta = getattr(picture, "meta", None)
    classification = getattr(meta, "classification", None) if meta is not None else None
    predictions = getattr(classification, "predictions", None) if classification is not None else None
    if predictions:
        return [predictions[0].class_name]

    classes: List[str] = []
    for ann in getattr(picture, "annotations", []):
        predicted = getattr(ann, "predicted_classes", None)
        if predicted:
            classes.append(predicted[0].class_name)
    return classes