"""Settle what a figure leaves on the page once the scan has read the figure.

Two answers come back from one look at the page: the text around the figure that
is only its furniture, which is removed, and the footnote markers its labels
carry, which are recorded on the picture.

The parse lifts what it can off a picture and files the fragments as text. Where
it files them under the picture, the graft removes them with the picture's
children. Where it files them under the page body instead — which it does
page to page with no pattern — they survive, and the document then holds a
chart's axis printed as forty separate records beside the reading that
supersedes them. Indexed for retrieval, each one carries the chart's heading, so
a gridline labelled 3.50 becomes a chunk that answers a question about the
weighted average risk rating with a number off the ruler.

The candidates are the plain text items overlapping a figure the scan read. That
is a bound on what can be removed rather than a decision about it: a heading, a
caption, a footnote or a page mark is never a candidate, whatever it overlaps,
and neither is a page the scan returned no figure for.

Deciding among the candidates is left to the agent, which is shown the page.
Nothing here inspects the text.

The markers are asked for on every page the scan returned a figure on, whether or
not that page has candidate text. A chart printing "(1,2)" on its title leaves
nothing behind to sweep, and it still carries a qualification the description
alone does not state.
"""

from __future__ import annotations

import asyncio
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

from docling_core.types.doc.document import DoclingDocument, PictureItem, PictureMeta, TextItem
from docling_core.types.doc.labels import DocItemLabel
from loguru import logger

from quber.agents.figure_correction import Figure, FigureCorrector, Fragment, Marker, Note
from quber.core.figures.geometry import NormBox, norm_box, prov_box
from quber.core.figures.graft import FOOTNOTE_MARKS_FIELD, FOOTNOTES_FIELD
from quber.core.figures.models import FigureRecord, PageScan, RemovedText
from quber.files.pdf import render_page

#: Page rasterization for the correction agent. It judges a fragment by where it
#: sits on the page, which is the same reading the table pipeline renders for.
CORRECTION_DPI = 200


async def correct_figures(
    document: DoclingDocument,
    scans: Sequence[PageScan],
    page_dims: Dict[int, Tuple[float, float]],
    source: Path,
    finder: Optional[FigureCorrector] = None,
) -> List[RemovedText]:
    """Sweep each page's figure furniture from `document` and record its markers.

    Each removal is returned with its text and its box, because a completeness
    check reading the source PDF cannot otherwise tell this from a silent loss:
    both leave printed text the document does not hold.

    Mutates `document` in place. Does nothing when the sweep is switched off, when
    no page carries a figure, or when a page's call fails.
    """
    if finder is None:
        return []

    work = [
        (scan, figures, candidates(document, scan, page_dims))
        for scan in scans
        for figures in [[figure for figure in scan.figures if figure.text]]
        if figures
    ]
    if not work:
        return []

    with tempfile.TemporaryDirectory(prefix="quber-correction-") as tmp:
        rendered = await asyncio.gather(
            *(
                asyncio.to_thread(
                    render_page, source, scan.page, CORRECTION_DPI, Path(tmp) / f"page-{scan.page:04d}.png"
                )
                for scan, _f, _c in work
            )
        )
        verdicts = await asyncio.gather(
            *(
                finder.correct_figure(
                    Path(image).read_bytes(),
                    [Figure(index=i, description=f.text) for i, f in enumerate(figures)],
                    [Fragment(index=i, text=item.text) for i, item in enumerate(eligible)],
                )
                for (_scan, figures, eligible), (image, _w, _h) in zip(work, rendered, strict=True)
            )
        )

    doomed = []
    removed: List[RemovedText] = []
    for (scan, figures, eligible), correction in zip(work, verdicts, strict=True):
        width, height = page_dims.get(scan.page, (612.0, 792.0))
        _record_markers(document, figures, correction.markers, correction.notes, scan.page)
        marked = [(eligible[f.index], f.reason) for f in correction.furniture]
        if marked:
            logger.info(
                "Figure correction: page {} removing {} of {} fragment(s) the scan's figures already read",
                scan.page,
                len(marked),
                len(eligible),
            )
        for item, reason in marked:
            doomed.append(item)
            removed.append(
                RemovedText(
                    page=scan.page,
                    text=item.text,
                    box=prov_box(item, width, height),
                    reason=reason,
                )
            )

    if not doomed:
        return []
    # Boxes are read before the delete: removing an item renumbers the ones after
    # it, and a reference resolved afterwards would name a different element.
    document.delete_items(node_items=doomed)
    return removed


def candidates(
    document: DoclingDocument, scan: PageScan, page_dims: Dict[int, Tuple[float, float]]
) -> List[TextItem]:
    """The plain text items on `scan`'s page that overlap a figure it returned."""
    boxes = [box for box in (norm_box(chart.box) for chart in scan.figures) if box is not None]
    if not boxes:
        return []
    width, height = page_dims.get(scan.page, (612.0, 792.0))
    found: List[TextItem] = []
    for item in document.texts:
        if item.label != DocItemLabel.TEXT or not (item.text or "").strip():
            continue
        if not item.prov or item.prov[0].page_no != scan.page:
            continue
        # A fragment already filed under a picture is removed with the picture's
        # children, so it is not put to the agent as well.
        if isinstance(item.parent.resolve(document) if item.parent else None, PictureItem):
            continue
        box = prov_box(item, width, height)
        if box is not None and any(_overlaps(figure, box) for figure in boxes):
            found.append(item)
    return found


def _record_markers(
    document: DoclingDocument,
    figures: Sequence[FigureRecord],
    markers: Sequence[Marker],
    notes: Sequence[Note],
    page: int,
) -> None:
    """Write each figure's markers and printed notes onto its picture.

    A figure the graft never placed has no picture to carry them. The run already
    reports that figure as unplaced, so the markers are logged and dropped rather
    than reported a second time.
    """
    if not markers and not notes:
        return
    by_ref = {picture.self_ref: picture for picture in document.pictures}
    for index, figure in enumerate(figures):
        mine = [m for m in markers if m.figure == index]
        read = [n for n in notes if n.figure == index]
        if not mine and not read:
            continue
        picture = by_ref.get(figure.picture_ref or "")
        if picture is None:
            logger.warning(
                "Figure markers: page {} figure {} carries {} marker(s) and sits on no picture",
                page,
                index,
                len(mine),
            )
            continue
        base = picture.meta or PictureMeta()
        picture.meta = base.model_copy(
            update={
                FOOTNOTE_MARKS_FIELD: [{"marker": m.marker, "kind": m.kind, "label": m.label} for m in mine],
                FOOTNOTES_FIELD: [{"marker": n.marker, "text": n.text} for n in read],
            }
        )
        logger.info(
            "Figure markers: page {} picture {} carries {} with {} note(s) read off the page",
            page,
            picture.self_ref,
            ", ".join(f"{m.marker} ({m.kind})" for m in mine) or "no markers",
            len(read),
        )


def _overlaps(figure: NormBox, text: NormBox) -> bool:
    """Do the two boxes share any area at all?

    Any overlap, with no fraction to tune. Measured on three decks: every text
    item on a scanned page either sits well inside a figure or misses it
    entirely, and not one landed in between, so a threshold decided nothing and
    only added a constant fitted to whichever deck it was read off.
    """
    wide = min(figure[2], text[2]) - max(figure[0], text[0])
    tall = min(figure[3], text[3]) - max(figure[1], text[1])
    return wide > 0 and tall > 0