"""Detect and repair Camelot capture drops before structure correction.

Camelot occasionally drops printed values during cell assignment even though
its text extraction captured them: a totals row typeset on a raised baseline
splits into two row bands at the default grouping tolerance, and the numbers'
vertical center lands exactly on the band boundary, failing the strict
containment test on both sides — the row's label and dollar signs survive,
its numbers vanish from the grid.

The repair loop keeps authority deterministic and puts the model only where
judgement is needed:

1. DETECT (deterministic): numeric tokens present in the table region's text
   layer but absent from the grid. Same numeric normalization as the
   correction-stage grounding guard.
2. RECOMMEND (agent): the capture advisor sees the cropped table image, the
   extracted grid, and the dropped tokens with their printed lines, and
   proposes one bounded retry adjustment.
3. RETRY + ACCEPT (deterministic): re-run the region pass with the advised
   knobs. The retry replaces the original ONLY if it recovers every missing
   numeric and loses none the original had. Anything else keeps the original
   grid, so a wrong recommendation costs one extra Camelot pass and nothing
   else.

Every detection and every verdict is logged — a capture drop on a financial
document must never be silent.
"""

from __future__ import annotations

import asyncio
from pathlib import Path
from typing import List, Optional, Set

from loguru import logger

from quber.agents.capture_advisor import CaptureAdvisor
from quber.agents.completeness import page_words
from quber.core.extractors.camelot.acquire import CamelotCandidate, grid_to_markdown
from quber.core.extractors.camelot.correspondence.correction import numeric_keys
from quber.core.extractors.camelot.correspondence.geometry import (
    WordBox,
    bbox_to_top_left,
    crop_region_png,
    table_crop_box,
)
from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted

# How far outside the Camelot bbox (PDF points) a word may sit and still count
# as inside the table for drop detection. Tight on purpose: a generous pad
# would pull a neighbouring table's numbers into the missing set and flag
# healthy captures.
DETECT_PAD_PTS = 2.0


def dropped_numeric_keys(cells: List[List[str]], region_words: List[WordBox]) -> Set[str]:
    """Numeric tokens in the region's text layer that appear in no grid cell."""
    layer = numeric_keys(" ".join(w[4] for w in region_words))
    grid = numeric_keys(" ".join(c for row in cells for c in row))
    return layer - grid


def printed_lines(region_words: List[WordBox], missing: Set[str]) -> List[str]:
    """For each missing token, the full printed line it sits on — the evidence
    the advisor grounds its diagnosis in."""
    lines: List[str] = []
    for key in sorted(missing):
        for w in region_words:
            if numeric_keys(w[4]) == {key}:
                line = " ".join(x[4] for x in region_words if abs(x[1] - w[1]) < 3)
                lines.append(f"'{key}' on printed line: {line}")
                break
        else:
            lines.append(f"'{key}' (word not isolated on a line)")
    return lines


async def repair_capture(
    cand: CamelotCandidate,
    source: str,
    page: int,
    page_image: Optional[Path],
    dpi: int,
    ordinal: int,
    advisor: Optional[CaptureAdvisor],
) -> CamelotCandidate:
    """Return `cand`, or an advised retry that provably captures more.

    No-op when the advisor is off, the candidate has no bbox, or no drop is
    detected. The retry is scoped to the candidate's own bbox — the exact
    frame the drop was measured in.
    """
    if advisor is None or cand.bbox is None or page_image is None:
        return cand
    _pw, page_h, words = await asyncio.to_thread(page_words, Path(source), page)
    left, top, right, bottom = bbox_to_top_left(cand.bbox, page_h, pad=DETECT_PAD_PTS)
    region_words = [w for w in words if w[0] >= left and w[2] <= right and w[1] >= top and w[3] <= bottom]
    missing = dropped_numeric_keys(cand.cells, region_words)
    if not missing:
        return cand
    logger.warning(
        "page {}: capture drop detected — {} numeric token(s) in the table's text layer "
        "missing from the grid: {}",
        page,
        len(missing),
        sorted(missing),
    )

    crop = await asyncio.to_thread(crop_region_png, page_image, table_crop_box(cand.bbox, page_h), dpi)
    advice = await advisor.recommend(
        crop, grid_to_markdown(cand.cells), printed_lines(region_words, missing), page
    )
    if advice is None:
        return cand
    logger.info(
        "page {}: capture advisor recommends flavor={} row_tol={} column_tol={}: {}",
        page,
        advice.flavor,
        advice.row_tol,
        advice.column_tol,
        advice.diagnosis,
    )

    x1, y1, x2, y2 = cand.bbox
    area = f"{min(x1, x2):.1f},{max(y1, y2):.1f},{max(x1, x2):.1f},{min(y1, y2):.1f}"
    try:
        retry = await asyncio.to_thread(
            camelot_targeted,
            source,
            page,
            area,
            ordinal,
            advice.flavor,
            advice.row_tol,
            advice.column_tol,
        )
    except Exception as exc:
        logger.warning("page {}: capture-repair retry failed; keeping original grid: {}", page, exc)
        return cand
    if retry is None:
        logger.warning("page {}: capture-repair retry found no grid; keeping original", page)
        return cand

    original = numeric_keys(" ".join(c for row in cand.cells for c in row))
    retried = numeric_keys(" ".join(c for row in retry.cells for c in row))
    recovered = missing & retried
    lost = original - retried
    if recovered == missing and not lost:
        logger.info(
            "page {}: capture repair ACCEPTED — recovered {} token(s), lost none",
            page,
            len(recovered),
        )
        return retry
    logger.warning(
        "page {}: capture repair REJECTED (recovered {}/{}, lost {}); keeping original grid",
        page,
        len(recovered),
        len(missing),
        len(lost),
    )
    return cand