"""Turn a table read off a page image into an extracted table.

A page is nominated partly for its tables: `table_verdicts` in nominate.py
counts the source PDF's text-layer words under each table's box against the
cells the table holds, and a table with too few was read off the page image
rather than out of a text layer. A table read off the image never reached the
Set-of-Mark and Camelot engine either, because that engine extracts from a text
layer and there was none.
What stands in the document for such a table is the parse's own reading, with
nothing having checked it since.

The scan replaces that reading. It returns the table as a table — a grid with a
box on every cell — so the grid goes straight into the same correction and
grounding steps every other table in the document goes through, and what comes
back is an `ExtractedTable` filled the way every other one is filled.

Replacing rather than keeping both is deliberate. The parse's reading of these
tables is not merely rougher, it is wrong in ways that change values: a currency
symbol read as a letter, a thousands separator read as a decimal point. A wrong
number carried beside a right one is how a wrong number survives.

The parse's reading is not discarded on the way out, though. Before the grid
reaches the correction step, every cell in it is checked against what the parse
read in the same place, and any cell the two disagree on, or the parse has no
cell for, is read off the printed page by an agent. The check is skipped when
the scan gave the table no box, and `crosscheck_cells` passes the grid through
unchecked when no reader is configured. The scan has its own failure mode — it returns a thousands
separator as a decimal point — and the parse is a second opinion on it that
costs nothing.

Only the tables that nominated the page are replaced. A page can print a table
with a text layer beside one without, and the first was read properly by the
table engine; overriding it with a scan would discard a vetted extraction for no
reason. A nominated table the scan returned nothing over is reported.
"""

from __future__ import annotations

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

from docling_core.types.doc.document import DocItem, DoclingDocument, PictureItem, TableItem
from docling_core.types.doc.labels import DocItemLabel
from loguru import logger

from quber.agents.cell_reader import CellReader
from quber.core.extractors.base import ExtractedTable
from quber.core.extractors.camelot.correspondence.geometry import (
    coverage_fraction,
    crop_region_png,
    table_crop_box,
)
from quber.core.extractors.set_of_mark.assemble import TableAssembly, assemble_table
from quber.core.extractors.set_of_mark.extent import apply_content_regions
from quber.core.figures.crosscheck import crosscheck_cells
from quber.core.figures.geometry import NormBox, norm_box, point_box, prov_box
from quber.core.figures.models import PageScan, ScannedTable
from quber.core.figures.nominate import FURNITURE_CLASSES, picture_classes
from quber.core.fusion.graft import READ_BY_FIELD
from quber.files.pdf import render_page

#: Minimum box overlap for a returned table and a parse table to be the same
#: region, measured both ways so a table drawn wider on either side registers.
#: The same threshold the figures are matched at.
MATCH_FRACTION = 0.20

#: Page rasterization for the correction and inspection agents. They read cell
#: text off the image, which is what the table pipeline renders at.
CAPTURE_DPI = 200


async def capture_tables(
    document: DoclingDocument,
    scans: List[PageScan],
    source: Path,
    page_dims: Dict[int, Tuple[float, float]],
    assembly: TableAssembly,
    reader: Optional[CellReader] = None,
    match_fraction: float = MATCH_FRACTION,
    orphans: bool = False,
) -> Tuple[List[ExtractedTable], List[str]]:
    """Extract every nominated table the scans read, in page order.

    Each record in `scans` is given the reference of the parse table its grid
    replaced and the address of the table it produced, so a value traces back to
    the element that holds it. Returns the tables and one error per nominated
    table no returned table covered.

    With `orphans`, a returned grid over a region the document holds nothing
    for — no table, no picture, no figure, and no text the parse read there —
    is captured too, through the same checks, with the cell check asking an
    agent about every cell because there is no parse reading to check against.
    A map's legend is the standing example: the page prints the values, the
    parse's text layer has nothing, and the scan's grid is the only reading
    anywhere. The caller gives such a table a home by inserting it.
    """
    by_ref = {t.self_ref: t for t in document.tables}
    work: List[Tuple[PageScan, ScannedTable, str, int]] = []
    errors: List[str] = []
    for scan in scans:
        width, height = page_dims.get(scan.page, (612.0, 792.0))
        ordinals = _page_ordinals(document, scan.page)
        # A returned table stands for one parse table. Claiming it keeps two
        # nominated tables in the same region from both taking the same grid.
        available = list(scan.tables)
        for ref in scan.table_refs:
            table_item = by_ref.get(ref)
            scanned = _best_match(table_item, available, width, height, match_fraction)
            if scanned is None:
                errors.append(
                    f"page {scan.page}: table {ref} was read off the page image and the scan "
                    "returned no table over it; what stands in the document is the parse's reading"
                )
                continue
            scanned.table_ref = ref
            available.remove(scanned)
            work.append((scan, scanned, ref, ordinals.get(ref, 1)))

        # A table the parse filed as a picture. The claim above runs over the
        # parse's tables and this region is not among them, so nothing took the
        # grid. It is still a table the page prints and the scan read, so the
        # picture stands in as its home and the grid travels the same path every
        # other scanned table travels — including the cell check, which asks an
        # agent about every cell when there is no parse reading to check against.
        ordinal = len(ordinals)
        for picture, scanned in _over_pictures(document, scan, available, width, height, match_fraction):
            scanned.picture_ref = picture.self_ref
            available.remove(scanned)
            ordinal += 1
            work.append((scan, scanned, picture.self_ref, ordinal))

        if orphans:
            for scanned in _orphan_grids(document, scan, available, width, height, match_fraction):
                available.remove(scanned)
                ordinal += 1
                work.append((scan, scanned, "", ordinal))

    if not work:
        return [], errors

    with tempfile.TemporaryDirectory(prefix="quber-scan-table-") as tmp:
        pages = sorted({scan.page for scan, _s, _r, _o in work})
        rendered = await asyncio.gather(
            *(
                asyncio.to_thread(render_page, source, page, assembly.dpi, Path(tmp) / f"page-{page:04d}.png")
                for page in pages
            )
        )
        images = {page: image for page, (image, _w, _h) in zip(pages, rendered, strict=True)}
        tables = await asyncio.gather(
            *(
                _capture(
                    scanned,
                    by_ref.get(ref),
                    images[scan.page],
                    page_dims.get(scan.page, (612.0, 792.0)),
                    ordinal,
                    assembly,
                    reader,
                )
                for scan, scanned, ref, ordinal in work
            )
        )

    for (_scan, scanned, _ref, _ordinal), table in zip(work, tables, strict=True):
        scanned.table_id = table.table_id
    # The scan bounds a table generously, so its box can run past the last data
    # row. Record the true end of the content from the corrected last row.
    await asyncio.to_thread(apply_content_regions, list(tables), str(source))
    return list(tables), errors


async def _capture(
    scanned: ScannedTable,
    table_item: Optional[TableItem],
    page_image: Path,
    page_dims: Tuple[float, float],
    ordinal: int,
    assembly: TableAssembly,
    reader: Optional[CellReader],
) -> ExtractedTable:
    """One returned grid, checked against the parse, then corrected and grounded."""
    width, height = page_dims
    boxes = [[point_box(box, width, height) for box in row] for row in scanned.cell_boxes]
    bbox = point_box(scanned.box, width, height)
    logger.info(
        "Table capture: page {} reading a {}x{} table off the page image",
        scanned.page,
        len(scanned.cells),
        len(scanned.cells[0]) if scanned.cells else 0,
    )
    # Values are settled against the printed page before the structure
    # correction runs, so that step works on a grid whose values are agreed.
    cells = scanned.cells
    if bbox is not None:
        crop = await asyncio.to_thread(
            crop_region_png, page_image, table_crop_box(bbox, height), assembly.dpi
        )
        cells, _read = await crosscheck_cells(scanned, table_item, page_dims, crop, reader)
    return await assemble_table(
        assembly,
        page=scanned.page,
        ordinal=ordinal,
        page_image=page_image,
        page_dims=page_dims,
        cells=cells,
        cell_boxes=boxes,
        bbox=bbox,
        som_region=norm_box(scanned.box),
        kind="image_table",
        # The page prints no text under this table — that is what made it one of
        # ours. So there is no text layer to check the corrected numbers against,
        # and checking them against the scan's own grid would reject every repair
        # of what the scan misread.
        ground_values=False,
    )


def _over_pictures(
    document: DoclingDocument,
    scan: PageScan,
    available: List[ScannedTable],
    width: float,
    height: float,
    match_fraction: float,
) -> List[Tuple[PictureItem, ScannedTable]]:
    """Each unclaimed grid paired with the picture it was printed over.

    A picture stands in for a table only where the parse filed the region as a
    picture and as nothing else. Four things disqualify it.

    Every class the parse gave it is page furniture (`FURNITURE_CLASSES`, a
    logo or an icon), so there is nothing inside it to read.

    A reader already read the region, and the picture's metadata records which
    under `READ_BY_FIELD`. Reading it again would state the same table twice.

    A parse table over the same region means the document already holds that
    table, read from a text layer by the table engine. The scan read it too, and
    that reading is ignored for the reason the module's own rule gives: a vetted
    extraction is not replaced by a scan of the same region. Adding it beside the
    table would state the same figures twice.

    A figure over the same region means the scan called it a picture, and the
    graft attaches that reading to the picture as a description. Taking it as a
    table as well would put one thing in the document twice again.
    """
    figures = [box for box in (norm_box(chart.box) for chart in scan.figures) if box is not None]
    parse_tables = [
        box
        for box in (
            prov_box(table, width, height)
            for table in document.tables
            if table.prov and table.prov[0].page_no == scan.page
        )
        if box is not None
    ]
    pairs: List[Tuple[PictureItem, ScannedTable]] = []
    for picture in document.pictures:
        if not picture.prov or picture.prov[0].page_no != scan.page:
            continue
        if picture_classes(picture) and all(c in FURNITURE_CLASSES for c in picture_classes(picture)):
            continue
        # A reader has already read this region and the document records which.
        # Reading it again would state the same table twice.
        if getattr(picture.meta, READ_BY_FIELD, None):
            continue
        target = prov_box(picture, width, height)
        if target is None:
            continue
        if any(_covers(other, target, match_fraction) for other in figures + parse_tables):
            continue
        remaining = [t for t in available if all(t is not p for _pic, p in pairs)]
        scanned = _best_match(picture, remaining, width, height, match_fraction)
        if scanned is not None:
            pairs.append((picture, scanned))
    return pairs


def _covers(box: NormBox, target: NormBox, match_fraction: float) -> bool:
    """Do the two boxes overlap enough, either way round, to be the same region?"""
    return max(coverage_fraction(box, target), coverage_fraction(target, box)) >= match_fraction


#: Labels that do not disqualify a region from being empty of parse text. A
#: heading floats over a region without holding its content, and page furniture
#: belongs to the page, not the region.
_SPINE_LABELS = frozenset(
    (DocItemLabel.SECTION_HEADER, DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER, DocItemLabel.TITLE)
)


def _orphan_grids(
    document: DoclingDocument,
    scan: PageScan,
    available: List[ScannedTable],
    width: float,
    height: float,
    match_fraction: float,
) -> List[ScannedTable]:
    """Each unclaimed grid over a region the document holds nothing for.

    Nothing means nothing: no parse table (the table engine read a text layer
    there), no figure (the graft attaches that reading to a picture), no
    picture (the picture tier above homes those), and no content text the
    parse read inside the region. That last condition is the line between a
    map's legend — printed values the parse's text layer never captured, where
    the scan's grid is the only reading anywhere — and a stat panel, whose
    strings the parse does hold and whose binding is the block grouping's job.
    Inserting a grid over a region the parse read would state its content
    twice.
    """
    figures = [box for box in (norm_box(chart.box) for chart in scan.figures) if box is not None]
    parse_tables = [
        box
        for box in (
            prov_box(table, width, height)
            for table in document.tables
            if table.prov and table.prov[0].page_no == scan.page
        )
        if box is not None
    ]
    pictures = [
        box
        for box in (
            prov_box(picture, width, height)
            for picture in document.pictures
            if picture.prov and picture.prov[0].page_no == scan.page
        )
        if box is not None
    ]
    text_centers = []
    for item in document.texts:
        if item.label in _SPINE_LABELS or not (item.text or "").strip():
            continue
        if not item.prov or item.prov[0].page_no != scan.page:
            continue
        box = prov_box(item, width, height)
        if box is not None:
            text_centers.append(((box[0] + box[2]) / 2, (box[1] + box[3]) / 2))

    found: List[ScannedTable] = []
    for scanned in available:
        if not scanned.cells:
            continue
        box = norm_box(scanned.box)
        if box is None:
            continue
        if any(_covers(other, box, match_fraction) for other in figures + parse_tables + pictures):
            continue
        if any(box[0] <= cx <= box[2] and box[1] <= cy <= box[3] for cx, cy in text_centers):
            continue
        found.append(scanned)
    return found


def _page_ordinals(document: DoclingDocument, page: int) -> Dict[str, int]:
    """Each table's position among the tables its page prints, counting from one."""
    refs = [t.self_ref for t in document.tables if t.prov and t.prov[0].page_no == page]
    return {ref: i + 1 for i, ref in enumerate(refs)}


def _best_match(
    item: Optional[DocItem],
    scanned: List[ScannedTable],
    width: float,
    height: float,
    match_fraction: float,
) -> Optional[ScannedTable]:
    """The returned table covering `item` best, or none past the threshold.

    `item` is whatever the parse holds over that region — a table it detected as
    one, or a picture it filed a table as.
    """
    if item is None:
        return None
    target = prov_box(item, width, height)
    if target is None:
        return None

    best: Optional[ScannedTable] = None
    best_cov = match_fraction
    for candidate in scanned:
        box = norm_box(candidate.box)
        if box is None:
            continue
        cov = max(coverage_fraction(box, target), coverage_fraction(target, box))
        if cov >= best_cov:
            best, best_cov = candidate, cov
    return best