"""
Chunk-to-table correspondence: assign Camelot chunks to detected tables
by 2D bbox overlap, deterministically — no LLM reads the cell values.

Camelot's chunk boxes are precise PDF geometry; the detector's boxes are
rough, so each detected table takes the chunk that covers it most
(argmax), not every chunk that clips its band — that ignores the spill a
rough boundary causes. Because the overlap is 2D this covers both
convergence axes (vertically stacked tables Camelot ran together, and
side-by-side tables it fused column-wise).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from quber.agents.detector import DetectedTable
from quber.core.extractors.camelot.acquire import CamelotCandidate
from quber.core.extractors.camelot.correspondence.geometry import (
    camelot_bbox_to_norm,
    coverage_fraction,
)

# A Camelot chunk is assigned to a detected table when their normalized
# rectangles overlap by at least this fraction of the detected table's
# area. The detector's boxes are rough, so the threshold is loose. A chunk
# that fills a table's box scores near 1.0 against it, and so does a chunk
# Camelot ran across several stacked or side-by-side tables, against each
# of them. An incidental sliver scores low and is rejected.
MATCH_FRACTION = 0.20


@dataclass(frozen=True)
class ChunkAssignment:
    """One flavor's argmax assignment of chunks to detected tables.

    `best_chunk` maps a detected-table index to the index of the chunk
    that covers it most. `tables_for_chunk` is the inversion, for every
    chunk: a chunk that is the best match for two-or-more tables is a
    combined chunk (kept whole, flagged, never split). `orphans` are the
    chunks that own no table, in the order the chunks were given — the
    only safe source for completeness extension, since pulling a chunk
    that owns another table would pollute this one.
    """

    best_chunk: Dict[int, int]
    tables_for_chunk: Dict[int, List[int]]
    orphans: List[CamelotCandidate]


def assign_chunks(
    detected: List[DetectedTable],
    remaining: List[int],
    chunks: List[CamelotCandidate],
    page_w_pts: float,
    page_h_pts: float,
) -> ChunkAssignment:
    """Assign each unresolved detected table (indices in `remaining`) to
    its best-overlap chunk. A table whose best coverage falls below
    MATCH_FRACTION, or whose detector box is missing, gets no assignment.
    `chunks` should already be in page order so `orphans` comes out in
    page order too.
    """
    chunk_norms: Dict[int, Optional[Tuple[float, float, float, float]]] = {
        ci: (camelot_bbox_to_norm(c.bbox, page_w_pts, page_h_pts) if c.bbox is not None else None)
        for ci, c in enumerate(chunks)
    }
    best_chunk: Dict[int, int] = {}
    for di in remaining:
        dbb = detected[di].bbox
        if dbb is None:
            continue
        best_ci, best_cov = None, 0.0
        for ci, cnorm in chunk_norms.items():
            if cnorm is None:
                continue
            cov = coverage_fraction(cnorm, dbb)
            if cov > best_cov:
                best_cov, best_ci = cov, ci
        if best_ci is not None and best_cov >= MATCH_FRACTION:
            best_chunk[di] = best_ci

    tables_for_chunk: Dict[int, List[int]] = {ci: [] for ci in range(len(chunks))}
    for di in sorted(best_chunk):
        tables_for_chunk[best_chunk[di]].append(di)

    orphans = [chunks[ci] for ci in range(len(chunks)) if not tables_for_chunk[ci]]
    return ChunkAssignment(best_chunk=best_chunk, tables_for_chunk=tables_for_chunk, orphans=orphans)


def chunk_top(cand: CamelotCandidate) -> float:
    """Sort key for top-to-bottom page order. Camelot bbox is in PDF
    points with a bottom-left origin, so a larger top-y sits higher on
    the page; we negate it. Chunks without a bbox sort last.
    """
    if cand.bbox is None:
        return float("inf")
    return -cand.bbox[3]


def assemble_cells(chunks: List[CamelotCandidate]) -> List[List[str]]:
    """Concatenate chunk cell grids in page order, values untouched. The
    structured grid is the canonical intermediate: assembly, the
    completeness audit and the fill all work on it, and markdown is
    rendered once at the output. Camelot stays the source of truth.
    """
    rows: List[List[str]] = []
    for c in chunks:
        rows.extend([list(r) for r in c.cells])
    return rows