"""Match the Set-of-Mark/Camelot tables against docling's tables and pictures.

One page at a time, every SoM table is overlapped against every docling table
and picture on that page. There are only a handful of tables per page, so the
match is all-pairs; no spatial index. Boxes from the three sources are first
normalized to one frame — 0..1 with the page top-left as origin — so they
compare directly:

- docling table / picture boxes: PDF points, bottom-left origin (flip y).
- SoM `som_region`: already normalized top-left (used as-is).
- Camelot `bbox` (fallback when a SoM table has no region): PDF points,
  bottom-left origin.

An edge is drawn when either box covers the other past `MATCH_FRACTION`, so a
large SoM region containing a small docling table and a small SoM table inside a
large docling table both register. SoM tables and docling tables are then grouped
into regions by connected overlap, and each region's cardinality is read off the
matrix in `models.py`.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

from docling_core.types.doc.base import CoordOrigin
from docling_core.types.doc.document import PictureItem, TableItem

from quber.core.extractors.base import ExtractedTable
from quber.core.extractors.camelot.correspondence.geometry import (
    camelot_bbox_to_norm,
    coverage_fraction,
)
from quber.core.fusion.models import RegionMatch
from quber.core.parsers.result import ParseResult, TableProvenance

#: The one knob: minimum box-overlap coverage for two tables to be considered the
#: same region. Carried over from the correspondence matcher's default.
MATCH_FRACTION = 0.20

NormBox = Tuple[float, float, float, float]


def match_tables(
    parse: ParseResult,
    som_tables: List[ExtractedTable],
    match_fraction: float = MATCH_FRACTION,
) -> List[RegionMatch]:
    """Match the SoM tables against the docling document, page by page.

    Returns one `RegionMatch` per region: every SoM table with a box lands in
    exactly one record (replace / som_merged / docling_undercount / image_table /
    chart / docling_miss), and every docling table with a provenance box that no
    SoM table matched yields a `som_miss` record. A SoM table with neither
    `som_region` nor `bbox`, and a docling table with no provenance, land in no
    record and are not reported.
    """
    document = parse.document
    page_dims = {p.page_no: (p.width, p.height) for p in parse.pages}
    prov_by_ref = parse.provenance_by_ref()

    pages = sorted(
        {t.page for t in som_tables}
        | {t.prov[0].page_no for t in document.tables if t.prov}
        | {p.prov[0].page_no for p in document.pictures if p.prov}
    )

    matches: List[RegionMatch] = []
    for page in pages:
        width, height = page_dims.get(page, (612.0, 792.0))
        matches.extend(
            _match_page(
                page,
                width,
                height,
                som_tables,
                document.tables,
                document.pictures,
                prov_by_ref,
                match_fraction,
            )
        )
    return matches


def _match_page(
    page: int,
    width: float,
    height: float,
    som_tables: List[ExtractedTable],
    docling_tables: List[TableItem],
    docling_pictures: List[PictureItem],
    prov_by_ref: Dict[str, TableProvenance],
    match_fraction: float,
) -> List[RegionMatch]:
    soms = [
        (i, _som_box(som_tables[i], width, height))
        for i in range(len(som_tables))
        if som_tables[i].page == page
    ]
    soms = [(i, b) for i, b in soms if b is not None]
    dtables = [
        (t.self_ref, _docling_box(t, width, height))
        for t in docling_tables
        if t.prov and t.prov[0].page_no == page
    ]
    dtables = [(ref, b) for ref, b in dtables if b is not None]
    dpics = [
        (p.self_ref, _docling_box(p, width, height), _picture_classes(p))
        for p in docling_pictures
        if p.prov and p.prov[0].page_no == page
    ]
    dpics = [(ref, b, c) for ref, b, c in dpics if b is not None]

    # SoM table -> docling tables it overlaps, and the best coverage seen.
    s_to_d: Dict[int, List[str]] = {i: [] for i, _ in soms}
    d_to_s: Dict[str, List[int]] = {ref: [] for ref, _ in dtables}
    best_cov: Dict[int, float] = {i: 0.0 for i, _ in soms}
    for i, sbox in soms:
        for ref, dbox in dtables:
            cov = max(coverage_fraction(sbox, dbox), coverage_fraction(dbox, sbox))
            if cov >= match_fraction:
                s_to_d[i].append(ref)
                d_to_s[ref].append(i)
                best_cov[i] = max(best_cov[i], cov)

    # Group SoM tables and docling tables that overlap into regions.
    components = _connected_regions([i for i, _ in soms], [ref for ref, _ in dtables], s_to_d)

    matches: List[RegionMatch] = []
    matched_soms: set[int] = set()
    matched_dtables: set[str] = set()
    for som_group, dtable_group in components:
        matched_soms.update(som_group)
        matched_dtables.update(dtable_group)
        matches.append(_classify_region(page, som_group, dtable_group, prov_by_ref, som_tables, best_cov))

    # SoM tables with no docling-table overlap: chart (over a picture) or a miss
    # docling did not see at all.
    for i, sbox in soms:
        if i in matched_soms:
            continue
        pic_refs: List[str] = []
        pic_classes: List[str] = []
        cov = 0.0
        for ref, pbox, classes in dpics:
            c = max(coverage_fraction(sbox, pbox), coverage_fraction(pbox, sbox))
            if c >= match_fraction:
                pic_refs.append(ref)
                pic_classes.extend(classes)
                cov = max(cov, c)
        if pic_refs:
            matches.append(
                RegionMatch(
                    page=page,
                    kind="chart",
                    som_indices=[i],
                    docling_picture_refs=pic_refs,
                    picture_classes=pic_classes,
                    overlap=cov,
                    detail="SoM table over a docling picture, no docling table",
                )
            )
        else:
            matches.append(
                RegionMatch(
                    page=page,
                    kind="docling_miss",
                    som_indices=[i],
                    detail="SoM table with no overlapping docling table or picture",
                )
            )

    # Docling tables no SoM table matched: SoM missed a table docling found.
    for ref, _ in dtables:
        if ref in matched_dtables:
            continue
        matches.append(
            RegionMatch(
                page=page,
                kind="som_miss",
                docling_table_refs=[ref],
                detail="docling table with no overlapping SoM table",
            )
        )

    return matches


def _classify_region(
    page: int,
    som_group: List[int],
    dtable_group: List[str],
    prov_by_ref: Dict[str, TableProvenance],
    som_tables: List[ExtractedTable],
    best_cov: Dict[int, float],
) -> RegionMatch:
    """Read one region's cardinality off the matrix."""
    c = len(som_group)
    d = len(dtable_group)
    cov = max((best_cov.get(i, 0.0) for i in som_group), default=0.0)

    # A 1:1 region where docling read the table off the image and Camelot found
    # nothing is a table rendered as an image, not a plain replace.
    if c == 1 and d == 1:
        prov = prov_by_ref.get(dtable_group[0])
        som = som_tables[som_group[0]]
        no_text_layer = prov is not None and prov.verdict in ("ocr", "empty")
        camelot_empty = not (som.markdown or "").strip() or som.camelot_accuracy <= 0.0
        if no_text_layer and camelot_empty:
            return RegionMatch(
                page=page,
                kind="image_table",
                som_indices=som_group,
                docling_table_refs=dtable_group,
                overlap=cov,
                detail=f"docling table read by OCR (verdict={prov.verdict if prov else 'n/a'}), Camelot empty",
            )

    if d > c:
        kind: str = "som_merged"
        detail = f"docling found {d} tables, SoM {c}: SoM merged stacked tables"
    elif c > d:
        kind = "docling_undercount"
        detail = f"SoM found {c} tables, docling {d}: docling dropped/merged a complex table"
    else:
        kind = "replace"
        detail = f"{c} SoM table(s) agree with {d} docling table(s)"

    return RegionMatch(
        page=page,
        kind=kind,
        som_indices=som_group,  # type: ignore[arg-type]
        docling_table_refs=dtable_group,
        overlap=cov,
        detail=detail,
    )


def _connected_regions(
    som_ids: List[int], dtable_refs: List[str], s_to_d: Dict[int, List[str]]
) -> List[Tuple[List[int], List[str]]]:
    """Group SoM tables and docling tables into connected overlap regions.

    Only SoM tables that overlap at least one docling table (and the docling
    tables they reach) form regions here; unmatched SoM tables and unmatched
    docling tables are handled by the caller.
    """
    parent: Dict[str, str] = {}

    def key_s(i: int) -> str:
        return f"s{i}"

    def key_d(ref: str) -> str:
        return f"d{ref}"

    def find(x: str) -> str:
        parent.setdefault(x, x)
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: str, b: str) -> None:
        parent.setdefault(a, a)
        parent.setdefault(b, b)
        parent[find(a)] = find(b)

    for i, refs in s_to_d.items():
        for ref in refs:
            union(key_s(i), key_d(ref))

    groups: Dict[str, Tuple[List[int], List[str]]] = {}
    for i in som_ids:
        if not s_to_d.get(i):
            continue
        root = find(key_s(i))
        groups.setdefault(root, ([], []))[0].append(i)
    for ref in dtable_refs:
        node = key_d(ref)
        if node not in parent:
            continue
        root = find(node)
        groups.setdefault(root, ([], []))[1].append(ref)

    return [g for g in groups.values() if g[0]]


def _som_box(table: ExtractedTable, width: float, height: float) -> Optional[NormBox]:
    """A SoM table's box in the normalized top-left frame.

    Prefers the SoM region (already normalized); falls back to Camelot's bbox
    (PDF points, bottom-left) when the region is absent.
    """
    if table.som_region is not None:
        return table.som_region
    if table.bbox is not None:
        return camelot_bbox_to_norm(table.bbox, width, height)
    return None


def _docling_box(item: TableItem | PictureItem, width: float, height: float) -> Optional[NormBox]:
    """A docling item's provenance bbox in the normalized top-left frame."""
    if not item.prov:
        return None
    bbox = item.prov[0].bbox
    if bbox.coord_origin == CoordOrigin.TOPLEFT:
        return (bbox.l / width, bbox.t / height, bbox.r / width, bbox.b / height)
    return camelot_bbox_to_norm((bbox.l, bbox.b, bbox.r, bbox.t), width, height)


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

    Predictions are ordered by descending confidence, so the first entry is
    docling's call (e.g. `bar_chart`). Reads the current `meta.classification`
    field, falling back to the deprecated `annotations` list for older output.
    """
    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