"""Vision+grid (Set-of-Mark) table-isolation approaches for the QUE-245 spike.

All approaches are Camelot-independent: they read the rendered page (with a
labeled grid overlaid) and report discrete grid IDs, which deterministic
geometry maps back to coordinates and snaps to the text layer. They differ
only in what the vision model is asked to bound:

  A "data-grid"   — the current production prompt: bound the data table
                    (header + aligned data cells); ignore page furniture.
  B "anatomy"     — bound the FULL table: title/caption + units line +
                    every column-header row (incl. spanning headers) + the
                    row-label stub column + all body/subtotal/total rows +
                    footnotes attached to the table. Separate stacked
                    tables; keep a single gapped table whole.

The grid readout, geometry, and text-layer tightening are shared, so an
A-vs-B difference is attributable to the instruction, not the machinery.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import List, Optional

from pydantic import BaseModel, Field

from quber.agents.grid_locator import (
    GridFlag,
    GridFlagResult,
    LocatedTable,
    build_prompt as build_data_grid_prompt,
    column_labels,
    flags_to_located,
    overlay_grid,
)

# --- Prompt A: current production "data-grid" prompt (re-exported) ------------

DATA_GRID_PROMPT = build_data_grid_prompt


# --- Prompt B: full-anatomy prompt -------------------------------------------


def anatomy_prompt(rows: int, cols: int) -> str:
    last = column_labels(cols)[-1]
    return f"""\
The page image has a grid overlaid: {rows} numbered horizontal rows (1 at the
top to {rows} at the bottom, labeled on BOTH side margins) and {cols} lettered
columns (A at the left to {last} at the right, labeled on the top and bottom
margins), with thin gridlines.

WHAT COUNTS AS A TABLE: a block where values are arranged in a GRID — the same
two or more columns line up vertically across multiple rows. These are tables.
The following are NOT tables, even when they contain dollar amounts or numbers,
and must never be flagged: a bulleted or numbered list; a single column of
figures; a heading followed by prose sentences; a paragraph. If the numbers do
not line up into repeating vertical columns, it is not a table.

Find every TABLE on the page and report the grid region that bounds each one
COMPLETELY. Bound the WHOLE table, including all of its parts:

- TITLE / CAPTION: the table's heading line(s) and any units line directly
  above the columns (e.g. "Condensed Consolidated Statements of Income",
  "(In millions)", "(UNAUDITED)"). Include them in the region.
- COLUMN HEADERS: every header row, including multi-level / spanning headers
  (a label like "Three Months Ended" that sits above several sub-columns, with
  the period and year rows beneath it). Include all header levels.
- ROW LABELS (the stub column): the left-hand column of row names. The region's
  LEFT edge must include it — do not start at the first number.
- BODY: every data row, including subtotal and total rows. Never cut a table
  off partway; blank lines inside a table do NOT end it.
- FOOTNOTES: short notes tied to the table by a marker (e.g. "(1)", "(2)")
  sitting just below it. Include them inside that table's region. Footnotes are
  NEVER a table of their own — never report them as a separate table.

Counting tables — judge by what you see, not by repeated text:
- Two tables stacked vertically are separate ONLY when each has its own
  title/caption and its own column-header block. Report them separately.
- One table may have section sub-headings inside it (e.g. "Cash flows from
  operating activities:", "Current assets:") and blank lines between sections.
  These do NOT start a new table — keep the whole thing as ONE region.
- Repeated "$" signs or a repeated subtotal style inside one table do NOT make
  a new table. When unsure, prefer ONE region over splitting.
Do NOT merge two genuinely separate tables, and do NOT split one table.

Exclude page furniture (running header timestamp, page number, document URL)
and free-standing prose paragraphs / lists that are not a table.

For each table, in reading order (top-to-bottom, then left-to-right), report:
- ordinal (1 = first)
- title (its visible caption, else a one-line summary of its columns)
- row_start: printed row number at the TOP of the title line (one row ABOVE the
  column headers when a title/caption is present) — include the title row
- row_end: printed row number at the BOTTOM-most part (last total, or footnote)
- col_start: printed column LETTER at the left edge, INCLUDING the row labels
- col_end: printed column LETTER at the right edge (last data column)

Read the printed labels off the grid — do not estimate coordinates.

Return ONLY a JSON object conforming to this schema (no prose, no code fences):

{json.dumps(GridFlagResult.model_json_schema(), indent=2)}
"""


# --- Flexible locator: same plumbing, swappable prompt/model -----------------


class FlexGridLocator:
    """A GridLocator whose prompt and model are parameters.

    Mirrors PydanticAIGridLocator's auth + decoding, but lets the spike swap
    the instruction (data-grid vs anatomy) and the model tier (haiku@temp0
    vs a higher tier with no temperature pin) without forking the class.
    """

    def __init__(
        self,
        system_prompt: str,
        model: str = "claude-haiku-4-5-20251001",
        temperature: Optional[float] = 0.0,
        rows: int = 24,
        cols: int = 12,
        snap: bool = False,
        pad_pts: float = 0.0,
    ) -> None:
        from pydantic_ai import Agent
        from pydantic_ai.settings import ModelSettings

        from quber.settings import get_settings

        llm = get_settings().llm
        auth = llm.anthropic_auth_token
        key = llm.anthropic_api_key
        if auth:
            from quber.agents._oauth_gate import make_oauth_anthropic_model

            anth = make_oauth_anthropic_model(model, auth)
        elif key:
            from pydantic_ai.models.anthropic import AnthropicModel
            from pydantic_ai.providers.anthropic import AnthropicProvider

            anth = AnthropicModel(model, provider=AnthropicProvider(api_key=key))
        else:
            raise RuntimeError("No ANTHROPIC_AUTH_TOKEN / ANTHROPIC_API_KEY available.")

        self.model = model
        self.rows = rows
        self.cols = cols
        self.snap = snap
        self.pad_pts = pad_pts
        self.system_prompt = system_prompt
        ms = None if temperature is None else ModelSettings(temperature=temperature)
        self.agent = Agent(anth, output_type=GridFlagResult, system_prompt=system_prompt, model_settings=ms)

    async def locate(self, page_image: Path, source: Path, page: int) -> List[LocatedTable]:
        import asyncio
        import tempfile

        from pydantic_ai import BinaryContent

        from quber.agents.completeness import page_words

        with tempfile.TemporaryDirectory(prefix="que245-grid-") as tmp:
            gridded = Path(tmp) / "gridded.png"
            overlay_grid(Path(page_image), gridded, self.rows, self.cols)
            image = BinaryContent(data=gridded.read_bytes(), media_type="image/png")
            result = await self.agent.run(["Flag the tables using the grid.", image])
            page_w, page_h, words = await asyncio.to_thread(page_words, Path(source), page)
        located = flags_to_located(result.output.tables, words, page_w, page_h, self.rows, self.cols)
        if self.pad_pts > 0.0:
            located = pad_boxes(located, page_w, page_h, self.pad_pts)
            located = declash_stacked(located, words, page_w, page_h)
        elif self.snap:
            located = declash_stacked(located, words, page_w, page_h)
            located = snap_titles(located, words, page_w, page_h)
        return located


# --- Text-layer title snap (precision refinement, not a boundary decision) ----


def _lines_above(words, top_pts: float, page_h: float):
    """Cluster words whose center sits above `top_pts` into visual lines."""
    above = [w for w in words if (w[1] + w[3]) / 2.0 < top_pts and str(w[4]).strip()]
    above.sort(key=lambda w: ((w[1] + w[3]) / 2.0))
    lines = []
    cur, cyc = [], None
    for w in above:
        yc = (w[1] + w[3]) / 2.0
        if cyc is None or abs(yc - cyc) <= 3.0:
            cur.append(w)
            cyc = yc if cyc is None else cyc
        else:
            lines.append(cur)
            cur, cyc = [w], yc
    if cur:
        lines.append(cur)
    return lines


_NUM_OK = set("0123456789.,()%-$ —–")


def _numeric_word_count(line) -> int:
    """Words in a line that read as numbers ($, %, parens, digits)."""
    n = 0
    for w in line:
        t = str(w[4]).strip()
        if any(c.isdigit() for c in t) and all(c in _NUM_OK for c in t):
            n += 1
    return n


def extend_to_title(
    region,
    words,
    page_w: float,
    page_h: float,
    ceiling_pts: float,
    max_gap_pts: float = 16.0,
):
    """Extend a box upward to capture a short caption/title line or two.

    Conservative by construction. Walking the lines directly above the box
    top, nearest first, a line is pulled in only when it is ALL of: close
    (gap <= max_gap_pts ~ one line), above the `ceiling` (the bottom of the
    table directly above — never climb into it), short (< 70% of the table
    width), and not a data row (fewer than 3 numeric words). The first line
    that fails any test stops the walk. A no-op when the model already
    bounded the title. Returns a normalized box.
    """
    x0, y0, x1, y1 = region
    cur_top = y0 * page_h
    table_w = (x1 - x0) * page_w
    for line in reversed(_lines_above(words, cur_top, page_h)):  # nearest first
        lx0 = min(w[0] for w in line)
        lx1 = max(w[2] for w in line)
        lty = min(w[1] for w in line)
        lby = max(w[3] for w in line)
        if lty <= ceiling_pts + 1.0:  # would enter the table above
            break
        gap = cur_top - lby
        if gap > max_gap_pts or gap < -2.0:
            break
        if (lx1 - lx0) > 0.70 * table_w:  # wide row -> data/prose, not a caption
            break
        if _numeric_word_count(line) >= 3:  # a data row, not a caption
            break
        cur_top = lty
    return (x0, cur_top / page_h, x1, y1)


def snap_titles(located, words, page_w: float, page_h: float):
    """Title-extend each box, never crossing into the table above it."""
    order = sorted(located, key=lambda t: min(t.region[1], t.region[3]))
    ceilings = {}
    prev_bottom = 0.0
    for t in order:
        ceilings[id(t)] = prev_bottom
        prev_bottom = max(t.region[1], t.region[3]) * page_h
    out = []
    for t in located:
        new_region = extend_to_title(t.region, words, page_w, page_h, ceilings[id(t)])
        out.append(t.model_copy(update={"region": new_region}))
    return out


def pad_boxes(located, page_w: float, page_h: float, pad_pts: float):
    """Grow every box by a fixed pad on all sides (clamped to the page).

    For a guidance box, slight over-inclusion is safe and clipping is not, so
    a dumb symmetric pad covers residual sub-grid-row edge misses without any
    content-dependent heuristic. Overlaps it creates between stacked tables
    are resolved by a following declash pass.
    """
    px, py = pad_pts / page_w, pad_pts / page_h
    out = []
    for t in located:
        x0, y0, x1, y1 = t.region
        r = (max(0.0, x0 - px), max(0.0, y0 - py), min(1.0, x1 + px), min(1.0, y1 + py))
        out.append(t.model_copy(update={"region": r}))
    return out


def declash_stacked(located, words, page_w: float, page_h: float, min_gap_pts: float = 2.0):
    """Remove vertical overlap between stacked table boxes at their seam.

    When an upper box's bottom runs into a lower box's top (grid quantization
    plus title-snap can pull the lower table's title into the upper box), pull
    the upper box's bottom UP to its own last row of text — the last word that
    sits above the lower box's top. The model already decided these are two
    tables; this only cleans the shared edge, it never merges or splits.
    """
    order = sorted(located, key=lambda t: min(t.region[1], t.region[3]))
    region_by_id = {id(t): list(t.region) for t in order}
    for upper, lower in zip(order, order[1:]):
        ur = region_by_id[id(upper)]
        u_top, u_bot = min(ur[1], ur[3]) * page_h, max(ur[1], ur[3]) * page_h
        seam = min(lower.region[1], lower.region[3]) * page_h
        if u_bot <= seam - min_gap_pts:
            continue  # no overlap
        ux0 = min(ur[0], ur[2]) * page_w
        ux1 = max(ur[0], ur[2]) * page_w
        sel = [
            w for w in words
            if u_top < (w[1] + w[3]) / 2.0 < seam - 0.5 and ux0 - 2 <= (w[0] + w[2]) / 2.0 <= ux1 + 2
        ]
        new_bot = min(max((w[3] for w in sel), default=seam), seam)
        region_by_id[id(upper)] = [ur[0], u_top / page_h, ur[2], new_bot / page_h]
    out = []
    for t in located:
        out.append(t.model_copy(update={"region": tuple(region_by_id[id(t)])}))
    return out


class _Region(BaseModel):
    region: tuple


class RawVisionLocator:
    """The continuous-coordinate baseline: the production table detector.

    Asks the vision model for each table's bbox directly (no grid, no text
    snap). This is the path the QUE-241 evidence found unreliable — boxes
    drift, clip, and merge. Included so the grid approaches are measured
    against the thing they would replace.
    """

    def __init__(self, model="claude-haiku-4-5-20251001", temperature: Optional[float] = 0.0):
        from quber.agents.detector import PydanticAITableDetector

        self.det = PydanticAITableDetector(model=model, temperature=temperature)

    async def locate(self, page_image: Path, source: Path, page: int):
        res = await self.det.detect(Path(page_image))
        return [_Region(region=tuple(t.bbox)) for t in res.tables if t.bbox is not None]


def make_locator(approach: str, model: str = "claude-haiku-4-5-20251001",
                 temperature: Optional[float] = 0.0, rows: int = 24, cols: int = 12,
                 snap: bool = False, pad_pts: float = 6.0):
    if approach == "raw":
        return RawVisionLocator(model=model, temperature=temperature)
    pad = 0.0
    if approach == "data-grid":
        prompt = build_data_grid_prompt(rows, cols)
    elif approach in ("anatomy", "anatomy-snap", "anatomy-pad"):
        prompt = anatomy_prompt(rows, cols)
        snap = snap or approach == "anatomy-snap"
        pad = pad_pts if approach == "anatomy-pad" else 0.0
    else:
        raise ValueError(f"unknown approach {approach!r}")
    return FlexGridLocator(prompt, model=model, temperature=temperature, rows=rows, cols=cols,
                           snap=snap, pad_pts=pad)