"""
Output-stage structure correction with a grounding guard.

The cropped table image is the arbiter of STRUCTURE (column/header
layout); Camelot's cells and the page text layer are the only sources of
VALUES. The LLM fixes Camelot's structural errors — a spurious empty
column from stream over-segmentation, a header split across rows, a
header misaligned from its values — and recovers any row Camelot
truncated from the text layer. When `ground_values` is on, a grounding
guard then rejects the correction if it introduced any value absent from
both Camelot's grid and the page text layer, and the caller falls back to
the deterministic grid markdown. The page-scan path in
`quber.core.figures.capture` turns the guard off for tables read off a page
image, so their corrections are not checked against any source.
"""

from __future__ import annotations

import asyncio
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Set, Tuple

from loguru import logger

from quber.agents.completeness import page_blocks, page_words
from quber.agents.llm_client import CellMerge, FootnoteDef, FootnoteMark, LLMClient
from quber.core.extractors.camelot.acquire import (
    grid_to_addressed_markdown,
    grid_to_markdown,
    is_content_empty,
)
from quber.core.extractors.camelot.correspondence.geometry import (
    CAPTION_PAD_PTS,
    bbox_to_top_left,
    crop_region_png,
    region_text_in_bbox,
    table_crop_box,
)
from quber.core.printed_text import is_printed, page_lines

__all__ = ["CAPTION_PAD_PTS", "StructureCorrection", "correct_structure", "numeric_keys", "printed_title"]


@dataclass(frozen=True)
class StructureCorrection:
    """Output-stage structure-correction result threaded onto make_extracted.

    `markdown` is the presentation render to emit -- the LLM's corrected
    markdown when accepted, else the deterministic grid render -- and
    `llm_corrected` records whether it differs from that grid render. The
    structured grid cells remain the value source of record; only the
    rendered presentation and the heading metadata change.
    """

    title: str
    # The sentence that introduces the table, copied off the page; empty when none.
    caption: str
    footnotes: List[FootnoteDef]
    markdown: str
    llm_corrected: bool
    # Leading column-header rows of the corrected markdown, counted off the
    # table image by the correction agent.
    header_rows: int
    units: str = ""
    footnote_refs: List[str] = field(default_factory=list)
    # Where each footnote marker sits, as the agent reported it: corrected
    # row/col plus the carrying cell's printed address. Validated into
    # LocatedMarkers at the grounding stage, not here.
    footnote_marks: List[FootnoteMark] = field(default_factory=list)
    # Cells the agent combined into one (split-symbol rejoin, header flatten),
    # each with the source values merged. Used to union source cell geometry
    # onto the merged cell.
    cell_merges: List[CellMerge] = field(default_factory=list)
    # The table region's text layer, already computed for the grounding guard.
    # Carried so the grounding stage can tell a printed-but-unmatched cell
    # from a label the correction added by convention.
    region_text: str = ""


async def correct_structure(
    cells: List[List[str]],
    page_image: Optional[Path],
    bbox: Optional[Tuple[float, float, float, float]],
    source: str,
    page: int,
    llm: LLMClient,
    correct_sem: asyncio.Semaphore,
    dpi: int,
    ground_values: bool = True,
) -> Optional[StructureCorrection]:
    """Vet a table's structure against its cropped image and ground the
    result in the page text layer.

    The image crop and the region text are scoped to this table's bbox. The
    agent also receives the whole page's text blocks, and the grounding guard
    accepts a number printed anywhere on the page, so a figure from an
    adjacent table on the same page passes the guard. Skipped when the page
    image or bbox is missing or the grid render is empty. Bounded by
    `correct_sem`. Any failure or rejection returns None so callers render the
    grid.

    `ground_values` checks every number in the corrected table against the grid
    and the page text layer, discarding the whole correction over one number
    neither carries. It is on for a table extracted from a text layer, where
    that layer is a source independent of the grid, so a number in neither has
    no origin.

    A caller whose table was read off a page image turns it off, because there
    the check has nothing independent to compare against — only the grid it
    asked to have corrected — so every repair fails it by construction. On a
    reviewed balance sheet every repair did: the correction restored the printed
    19,543,903 from a reading of 19.543.903 and was rejected for it, so the
    corrupted figure stood. Where the page prints no text, the agent reading the
    image is the only component that has seen the table.
    """
    if page_image is None or bbox is None:
        return None
    grid_md = grid_to_markdown(cells)
    if is_content_empty(grid_md):
        return None

    _pw, page_h_pts, words = await asyncio.to_thread(page_words, Path(source), page)
    region = bbox_to_top_left(bbox, page_h_pts)
    region_text = region_text_in_bbox(words, region)
    # The whole page as its text blocks, the units the layout keeps together:
    # a paragraph, a title printed over two lines, one panel of a slide. The
    # caption and title are copied from here, because the region's text layer
    # cuts a line in half where two tables sit side by side. What is copied is
    # then checked against the blocks and against the page's printed lines,
    # since either reading can hold a phrase the other splits.
    blocks = await asyncio.to_thread(page_blocks, Path(source), page)
    page_text = "\n\n".join(blocks)
    printed = _printed_text(blocks, words)
    # The image crop reaches a bit higher than the data box to include the
    # table's title/caption, which sits just above the grid. The grounding
    # guard below checks numbers against every word on the page, not only
    # `region_text`, so a neighbouring table's number on the same page is
    # accepted by it.
    crop_box = table_crop_box(bbox, page_h_pts)
    try:
        crop_png = await asyncio.to_thread(crop_region_png, page_image, crop_box, dpi)
    except Exception as exc:
        logger.error("page {}: table region crop failed: {}", page, exc)
        return None

    async with correct_sem:
        try:
            # The agent sees the grid inside a printed coordinate frame so its
            # merge reports can NAME source cells by address (read, not
            # counted). The guard and the llm_corrected comparison below use
            # the unlabeled render — the frame is reference only.
            correction = await llm.vet_structure(
                crop_png, grid_to_addressed_markdown(cells), region_text, page_text
            )
        except Exception as exc:
            logger.error("page {}: LLM structure vetting failed: {}", page, exc)
            return None
    if correction is None or not (correction.markdown and correction.markdown.strip()):
        return None
    for opening in correction.body_text:
        logger.info("page {}: running text below the table left out of it: {!r}", page, opening)

    # The coordinate frame must never leak into the corrected table. A leaked
    # frame would corrupt values downstream, so it rejects the correction.
    if _leaks_coordinate_frame(correction.markdown):
        logger.warning(
            "page {}: structure correction REJECTED (coordinate frame leaked into output); "
            "keeping grid markdown",
            page,
        )
        return None

    corrected = correction.markdown
    if ground_values:
        ungrounded = ungrounded_values(corrected, grid_md, " ".join(w[4] for w in words))
        if ungrounded:
            logger.warning(
                "page {}: structure correction REJECTED (ungrounded values {}); keeping grid markdown",
                page,
                sorted(ungrounded),
            )
            return None

    corrected = _blank_unsourced_placeholders(corrected, grid_md, region_text)

    # The caption and the title are copies of printed text or nothing. A value
    # the page does not print was composed, and a composed header line would
    # steer retrieval toward words the document never says.
    title = correction.title if is_printed(correction.title, printed) else ""
    caption = correction.caption if is_printed(correction.caption, printed) else ""
    for name, given, kept in (("title", correction.title, title), ("caption", correction.caption, caption)):
        if given and not kept:
            logger.info("page {}: {} is not printed on the page; dropped: {!r}", page, name, given[:80])

    return StructureCorrection(
        title=title,
        caption=caption,
        footnotes=list(correction.footnotes),
        markdown=corrected,
        llm_corrected=corrected != grid_md,
        units=correction.units,
        header_rows=correction.header_rows,
        footnote_refs=list(correction.footnote_refs),
        footnote_marks=list(correction.footnote_marks),
        # A merge whose joined cells were all blank names no printed mark at
        # all (CellMerge already drops blank pieces): not a provenance record.
        cell_merges=[m for m in correction.cell_merges if m.sources],
        region_text=region_text,
    )


_DASH_ONLY_RE = re.compile(r"^[—–‒―-]+$")
_SEPARATOR_CELL_RE = re.compile(r"^[-: ]+$")


#: Placed between blocks and between lines before a printed-text check. The
#: check ignores whitespace so a phrase split over two lines of one block still
#: matches; this character, which is not whitespace, keeps two neighbouring
#: blocks from reading as one.
_BOUNDARY = "\u00b6"


def _printed_text(blocks: List[str], words: List[Tuple[float, float, float, float, str]]) -> str:
    """Everything the page prints, read twice: as text blocks and as lines,
    each unit closed off so a phrase cannot match across two of them."""
    return _BOUNDARY.join(blocks) + _BOUNDARY + _BOUNDARY.join(page_lines(words))


async def printed_title(title: str, source: str, page: int) -> str:
    """`title` when the page prints it, else empty.

    For a table whose structure correction did not run, the only title on offer
    is the one the grid locator read off the page image. That reading is not
    checked anywhere else, and a composed name would steer retrieval toward
    words the document never says, so it passes the same check a copied title
    does."""
    if not title or not title.strip():
        return ""
    _pw, _ph, words = await asyncio.to_thread(page_words, Path(source), page)
    blocks = await asyncio.to_thread(page_blocks, Path(source), page)
    if is_printed(title, _printed_text(blocks, words)):
        return title
    logger.info("page {}: located title is not printed on the page; dropped: {!r}", page, title[:80])
    return ""


def _blank_unsourced_placeholders(markdown: str, grid_md: str, region_text: str) -> str:
    """Blank every dash-only cell whose dash appears nowhere in Camelot's grid
    or the table's text layer.

    The vet prompt forbids writing a placeholder into a printed-blank cell,
    but compliance is not exact: the agent still occasionally emits an em-dash
    for a cell the page shows empty. Such a dash names nothing on the page. A
    genuinely printed nil dash is present in the grid or the text layer and is
    always kept. Normalizing here makes the output deterministic either way —
    the same boundary pattern as `CellMerge.drop_blank_sources`. Separator
    lines are structural, never touched.
    """
    source = grid_md + "\n" + region_text
    lines: List[str] = []
    for line in (markdown or "").splitlines():
        stripped = line.strip()
        if stripped.startswith("|"):
            cells = stripped.strip("|").split("|")
            if not all(_SEPARATOR_CELL_RE.fullmatch(c) for c in cells):
                cells = [
                    "  " if _DASH_ONLY_RE.fullmatch(c.strip()) and c.strip() not in source else c
                    for c in cells
                ]
                line = "|" + "|".join(cells) + "|"
        lines.append(line)
    return "\n".join(lines)


_TAG_LEAK_RE = re.compile(r"\[[A-Z]+\d+\]")


def _leaks_coordinate_frame(markdown: str) -> bool:
    """True when the corrected table still carries an inline address tag
    (e.g. '[B3]'), which would corrupt cell text downstream."""
    return bool(_TAG_LEAK_RE.search(markdown or ""))


def ungrounded_values(corrected: str, grid_md: str, page_text: str) -> Set[str]:
    """Numbers in the corrected table that neither the grid nor the page carries.

    Every number in a corrected table must be present in the extracted grid or
    in the page's text layer. A number in neither was read off nothing, and the
    caller discards the whole correction rather than carry it.
    """
    return numeric_keys(corrected) - (numeric_keys(grid_md) | numeric_keys(page_text))


_NUMERIC_RE = re.compile(r"\d[\d,]*(?:\.\d+)?")


def numeric_keys(text: str) -> Set[str]:
    """Every number in `text`, normalized to bare digits (plus a decimal point)
    for presence grounding.

    Currency, percent, parentheses, commas and spacing are dropped, so the same
    figure matches however it was rendered: '$ 8,273.04', '8273.04' and
    '8,273.04%' all key to '8273.04'. Bare integers count too — so a percentage
    column written as plain '21' in the source still grounds an output cell
    rendered as '21%'. This is presence grounding (is the figure in the source),
    not literal-drift detection, so a value is never rejected merely because a
    symbol moved into or out of its cell during correction.
    """
    return {m.group(0).replace(",", "") for m in _NUMERIC_RE.finditer(text)}