"""Set-of-Mark vision-guided extraction as a pydantic-graph.

Topology (mirrors the GraphBuilder shape of the LLM pipeline; every map edge
pairs with its own closest-fork join):

    start -> locate_tables -> [map] capture_table -> (join) -> finalize -> end

- locate_tables: render-fed page images in, the grid locator names every table
  and its region; emits one ref per table.
- capture_table (mapped per table): region-constrained Camelot extracts the
  cell grid inside the table's box. `repair_capture` (camelot/recapture.py)
  retries Camelot on the capture advisor's advice when the grid dropped values
  the region's text layer holds. `assemble_table` (assemble.py) then runs the
  grounded structure correction (spans, merged-symbol columns, multi-level
  headers) without touching values, grounds every corrected cell, and has the
  status inspector (inspection.py) verify the statuses proposed for unboxed
  cells against the table image.
- finalize: sort to reading order, then record each table's `content_region`
  through `apply_content_regions` (extent.py).

Vision owns identity and region; Camelot owns the values inside that region.
The structure-correction LLM cleans structure under the no-number guard. The
capture advisor only proposes a Camelot retry adjustment, and the status
inspector only checks cell statuses against the image.

The deprecated orchestrators in llm/ and correspondence/
(`CamelotLLMTableExtractor`, `CamelotCorrespondenceExtractor`) are still
reachable through `quber table --engine`, and `quber validate` builds
`CamelotLLMTableExtractor`. The correspondence package's correction, geometry
and recovery modules are live dependencies of this pipeline.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from functools import cache
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from pydantic import BaseModel, ConfigDict
from pydantic_graph import GraphBuilder, StepContext, reduce_list_append
from pydantic_graph.graph_builder import Graph
from pydantic_graph.id_types import JoinID, NodeID

from quber.agents.capture_advisor import CaptureAdvisor
from quber.agents.grid_locator import GridLocator, LocatedTable
from quber.agents.llm_client import LLMClient
from quber.agents.status_inspector import StatusInspector
from quber.core.extractors.base import ExtractedTable, table_address
from quber.core.extractors.camelot.correspondence.correction import printed_title
from quber.core.extractors.camelot.correspondence.geometry import norm_bbox_to_table_area
from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted
from quber.core.extractors.camelot.recapture import repair_capture
from quber.core.extractors.set_of_mark.assemble import CamelotOrigin, TableAssembly, assemble_table
from quber.core.extractors.set_of_mark.extent import apply_content_regions


@dataclass
class SetOfMarkDeps:
    """Run-scoped dependencies and the pre-rendered page images."""

    locator: GridLocator
    llm: LLMClient
    source: Path
    page_images: Dict[int, Path]
    page_dims: Dict[int, Tuple[float, float]]
    dpi: int
    correct_sem: asyncio.Semaphore
    page_sem: asyncio.Semaphore
    # Recommends a Camelot retry when a capture drop is detected; None
    # disables the repair loop (backend `off`).
    advisor: Optional[CaptureAdvisor] = None
    # Verifies each unboxed cell's proposed status against the table image;
    # None disables inspection (backend `off`) and proposed statuses stand.
    inspector: Optional[StatusInspector] = None


@dataclass
class SetOfMarkState:
    """No cross-step mutable state; everything run-scoped rides on deps."""


class LocatedTableRef(BaseModel):
    """One located table threaded into the per-table map stage."""

    model_config = ConfigDict(arbitrary_types_allowed=True)
    page: int
    located: LocatedTable


class CapturedTable(BaseModel):
    """A captured table plus its page/ordinal, so finalize can reading-order it."""

    page: int
    ordinal: int
    table: ExtractedTable


async def locate_tables(ctx: StepContext[SetOfMarkState, SetOfMarkDeps, Path]) -> List[LocatedTableRef]:
    deps = ctx.deps

    async def per_page(page: int) -> List[LocatedTableRef]:
        async with deps.page_sem:
            located = await deps.locator.locate(deps.page_images[page], deps.source, page)
        return [LocatedTableRef(page=page, located=t) for t in sorted(located, key=lambda t: t.ordinal)]

    per_page_refs = await asyncio.gather(*(per_page(p) for p in sorted(deps.page_images)))
    return [ref for page_refs in per_page_refs for ref in page_refs]


async def capture_table(ctx: StepContext[SetOfMarkState, SetOfMarkDeps, LocatedTableRef]) -> CapturedTable:
    deps = ctx.deps
    ref = ctx.inputs
    page, located = ref.page, ref.located
    page_w, page_h = deps.page_dims[page]
    area = norm_bbox_to_table_area(located.region, page_w, page_h)
    try:
        cand = await asyncio.to_thread(camelot_targeted, str(deps.source), page, area, located.ordinal)
    except Exception:
        cand = None

    if cand is None:
        # Region produced no grid: keep the visual identity, empty body.
        table = ExtractedTable(
            table_id=table_address(deps.source, page, located.ordinal),
            title=await printed_title(located.title, str(deps.source), page),
            markdown="",
            page=page,
            source=str(deps.source),
            flavor="stream",
            som_region=located.region,
        )
        return CapturedTable(page=page, ordinal=located.ordinal, table=table)

    # Capture-drop repair: values Camelot's cell assignment dropped (present in
    # the region's text layer, absent from the grid) trigger an advised retry,
    # accepted only when it provably captures more. See camelot/recapture.py.
    cand = await repair_capture(
        cand, str(deps.source), page, deps.page_images[page], deps.dpi, located.ordinal, deps.advisor
    )

    table = await assemble_table(
        TableAssembly(deps.source, deps.llm, deps.correct_sem, deps.dpi, deps.inspector),
        page=page,
        ordinal=located.ordinal,
        page_image=deps.page_images[page],
        page_dims=(page_w, page_h),
        cells=cand.cells,
        cell_boxes=cand.cell_boxes,
        bbox=cand.bbox,
        title=located.title,
        som_region=located.region,
        camelot=CamelotOrigin(cand.bbox, cand.flavor, cand.accuracy),
    )
    return CapturedTable(page=page, ordinal=located.ordinal, table=table)


async def finalize(
    ctx: StepContext[SetOfMarkState, SetOfMarkDeps, List[CapturedTable]],
) -> List[ExtractedTable]:
    ordered = sorted(ctx.inputs, key=lambda c: (c.page, c.ordinal))
    tables = [c.table for c in ordered]
    # The locator and Camelot both bound a table generously, so its box can run
    # past the last data row (e.g. enclosing footnote lines below the grid).
    # Record the true end of the table content from its corrected last row.
    await asyncio.to_thread(apply_content_regions, tables, str(ctx.deps.source))
    return tables


@cache
def build_set_of_mark_graph() -> Graph[SetOfMarkState, SetOfMarkDeps, Path, List[ExtractedTable]]:
    """Build the vision-guided extraction graph (cached; immutable)."""
    g = GraphBuilder(
        name="set_of_mark.pipeline",
        state_type=SetOfMarkState,
        deps_type=SetOfMarkDeps,
        input_type=Path,
        output_type=List[ExtractedTable],
    )

    locate_step = g.step(locate_tables)
    capture_step = g.step(capture_table)
    finalize_step = g.step(finalize)

    j_captured = g.join(
        reduce_list_append,
        initial_factory=list,
        node_id="j_captured",
        preferred_parent_fork="closest",
    )

    g.add_edge(g.start_node, locate_step)
    g.add_mapping_edge(locate_step, capture_step, downstream_join_id=JoinID(NodeID("j_captured")))
    g.add_edge(capture_step, j_captured)
    g.add_edge(j_captured, finalize_step)
    g.add_edge(finalize_step, g.end_node)

    return g.build()