"""Split a Set-of-Mark region that fused two or more stacked tables into one.

The Set-of-Mark locator occasionally marks two vertically-stacked,
independently-titled tables as a single region. The fused region extracts with
correct values but a broken structure: the two tables' period labels, headers
and footnotes get crossed or lost. `split_table` is a post-pass over one
already-captured table that refines such a region back into the separate tables
it should have been.

It is self-contained: same type in, same type out. The expected table count and
the per-table boundaries are handed in by the caller (their provenance is a
parsed document's table geometry); this function never parses the document
itself. It reaches the source PDF only through `table.source`: to render the
page image for the vision probe and the structure correction, to re-extract
values with Camelot within a tighter box, and to read the page text layer the
structure correction grounds against.

Two independent count signals decide whether to split: the handed-in `count`
and a vision count probe run here over the region image. They must agree on more
than one table; on agreement of one, on any disagreement, or on any failure of
the re-extraction, the region is returned exactly as it came in. The cut between
sub-regions is deterministic geometry taken from the supplied boundaries, never
an estimated coordinate, and the sub-regions always tile the full original
region so a footnote sitting below a table's tight box is never dropped.
"""

from __future__ import annotations

import asyncio
import tempfile
from pathlib import Path
from typing import List, Optional, Sequence, Tuple

import fitz
from loguru import logger

from quber.agents.llm_client import LLMClient
from quber.core.extractors.base import ExtractedTable, MergedCellBox, grid_fingerprint, grounded_grid
from quber.core.extractors.camelot.acquire import grid_to_markdown
from quber.core.extractors.camelot.correspondence.correction import correct_structure
from quber.core.extractors.camelot.correspondence.geometry import (
    crop_region_png,
    norm_bbox_to_table_area,
)
from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted
from quber.core.extractors.set_of_mark.merge_grounding import (
    find_dropped_header_text,
    locate_markers,
    log_ungrounded,
    log_ungrounded_cells,
    resolve_corrected_grid,
    resolve_merges,
)

# A box normalized to 0..1 with the page top-left as origin (x1, y1, x2, y2);
# the same frame as ExtractedTable.som_region. Boundaries handed in must share
# this frame so the cut geometry lines up with the region being subdivided.
NormBox = Tuple[float, float, float, float]


async def split_table(
    table: ExtractedTable,
    count: int,
    boundaries: Sequence[NormBox],
    llm: LLMClient,
    *,
    dpi: int = 200,
    correct_sem: Optional[asyncio.Semaphore] = None,
) -> List[ExtractedTable]:
    """Refine one captured table into the separate tables it fused, or pass it through.

    `count` is the expected number of tables in the region and `boundaries` are
    their per-table boxes (one per expected table, same normalized top-left frame
    as `table.som_region`). The region is split into `count` tables only when the
    vision probe agrees the region holds more than one table and every sub-region
    re-extracts cleanly; otherwise the input table is returned as a one-element
    list, unchanged.
    """
    region = table.som_region
    # Nothing to refine unless more than one table is expected, every expected
    # table has a boundary to cut on, and we have the region plus a source to
    # re-read. Any of these missing means the region stands as captured.
    if count <= 1 or region is None or table.source is None or len(boundaries) != count:
        return [table]

    sub_regions = subdivide(region, boundaries)
    if sub_regions is None or len(sub_regions) != count:
        return [table]

    source = Path(table.source)
    sem = correct_sem or asyncio.Semaphore(1)

    # Only a region that passed the checks above reaches this point. Every such
    # region pays for a page render and a vision count call, even when the probe
    # disagrees and the region is kept. Camelot re-extracts and corrections run
    # sub-region by sub-region and are spent even when a later one fails and the
    # split is abandoned.
    with tempfile.TemporaryDirectory(prefix="quber-split-") as tmp:
        try:
            page_image, page_w, page_h = await asyncio.to_thread(
                render_page, source, table.page, dpi, Path(tmp)
            )
        except Exception as exc:
            logger.error("page {}: split render failed: {}", table.page, exc)
            return [table]

        vision_count = await vision_count_probe(llm, page_image, region, page_w, page_h, dpi)
        # The handed-in count and the vision probe must agree on the same N > 1.
        # Anything else keeps the region as Camelot already extracted it.
        if vision_count != count:
            return [table]

        out: List[ExtractedTable] = []
        for ordinal, sub in enumerate(sub_regions, start=1):
            captured = await reextract_subregion(
                sub, source, table.page, ordinal, llm, page_image, page_w, page_h, dpi, sem
            )
            if captured is None:
                # A sub-region with no recoverable grid would drop content if
                # emitted, so the whole split is abandoned and the fused region
                # is kept intact rather than shipped partial.
                logger.warning(
                    "page {}: sub-region {} of {} produced no grid; keeping fused region",
                    table.page,
                    ordinal,
                    count,
                )
                return [table]
            # A sub-table extends its parent's address with its own sub-index.
            if table.table_id:
                captured.table_id = f"{table.table_id}-s{ordinal}"
            out.append(captured)

        # A sub-table's crop pads only a little above its own grid, so the unit
        # caption sitting above the whole group is seen by the top sub-table and
        # missed by the ones below it. Carry the last caption seen downward: each
        # sub-table with no caption of its own inherits the one above it, and a
        # sub-table that captured its own caption keeps it and carries that one to
        # the sub-tables below. The fill only flows down, so a sub-table above the
        # first caption stays empty.
        carried = ""
        for sub in out:
            if sub.units:
                carried = sub.units
            elif carried:
                sub.units = carried
        return out


def subdivide(region: NormBox, boundaries: Sequence[NormBox]) -> Optional[List[NormBox]]:
    """Tile `region` into vertical bands, one per boundary, cut in the gaps between them.

    Boundaries are ordered top-to-bottom by their vertical center; the cut for
    each adjacent pair sits at the midpoint of the blank gap between them. Every
    band spans the region's full width, and the bands together cover the region
    from its top edge to its bottom edge exactly — so content between a tight
    boundary box and the next cut (a table's footnote) stays inside its band
    rather than being dropped. Returns None if a boundary center falls outside
    the region or the cuts would not increase top-to-bottom.
    """
    rx1, ry1, rx2, ry2 = region
    left, right = min(rx1, rx2), max(rx1, rx2)
    top, bottom = min(ry1, ry2), max(ry1, ry2)

    def center_y(box: NormBox) -> float:
        return (min(box[1], box[3]) + max(box[1], box[3])) / 2.0

    boxes = sorted(boundaries, key=center_y)
    for box in boxes:
        if not (top <= center_y(box) <= bottom):
            return None

    seams: List[float] = []
    for upper, lower in zip(boxes, boxes[1:], strict=False):
        upper_bottom = max(upper[1], upper[3])
        lower_top = min(lower[1], lower[3])
        seams.append((upper_bottom + lower_top) / 2.0)

    edges = [top, *seams, bottom]
    if any(edges[i] >= edges[i + 1] for i in range(len(edges) - 1)):
        return None
    return [(left, edges[i], right, edges[i + 1]) for i in range(len(edges) - 1)]


async def vision_count_probe(
    llm: LLMClient,
    page_image: Path,
    region: NormBox,
    page_w: float,
    page_h: float,
    dpi: int,
) -> int:
    """Ask vision how many independent tables are viewable inside the region.

    Crops the rendered page to the region and counts tables in just that crop,
    so a multi-table page does not inflate the count. Returns 0 on any failure,
    which reads as disagreement and keeps the region unsplit.
    """
    region_pts = (region[0] * page_w, region[1] * page_h, region[2] * page_w, region[3] * page_h)
    try:
        crop_png = await asyncio.to_thread(crop_region_png, page_image, region_pts, dpi)
        crop_path = page_image.parent / "split-probe.png"
        crop_path.write_bytes(crop_png)
        return await llm.count_tables(crop_path)
    except Exception as exc:
        logger.error("split vision count probe failed: {}", exc)
        return 0


async def reextract_subregion(
    sub_region: NormBox,
    source: Path,
    page: int,
    ordinal: int,
    llm: LLMClient,
    page_image: Path,
    page_w: float,
    page_h: float,
    dpi: int,
    correct_sem: asyncio.Semaphore,
) -> Optional[ExtractedTable]:
    """Re-read one sub-region from the source PDF and clean its structure.

    Region-constrained Camelot reads the source within the sub-region box (the
    only source of values), then the same structure-correction step the pipeline
    runs per table fixes the header layout and recovers the title, caption and
    footnotes from the page image. Returns None if Camelot finds no grid in the
    box, so the caller can abandon the split rather than emit an empty table.
    """
    area = norm_bbox_to_table_area(sub_region, page_w, page_h)
    try:
        cand = await asyncio.to_thread(camelot_targeted, str(source), page, area, ordinal)
    except Exception as exc:
        logger.error("page {}: sub-region {} Camelot re-extract raised: {}", page, ordinal, exc)
        return None
    if cand is None:
        return None

    correction = await correct_structure(
        cand.cells, page_image, cand.bbox, str(source), page, llm, correct_sem, dpi
    )
    cell_grid = grounded_grid(cand.cells, cand.cell_boxes, page_w, page_h)
    merged: List[MergedCellBox] = []
    dropped_text: List[str] = []
    corrected_grid = cell_grid
    if correction is not None:
        if correction.cell_merges:
            merged = resolve_merges(correction.cell_merges, cell_grid, correction.markdown)
            log_ungrounded(merged, page)
        corrected_grid = resolve_corrected_grid(
            correction.markdown, merged, cell_grid, correction.footnote_refs, correction.region_text
        )
        log_ungrounded_cells(corrected_grid, page)
        dropped_text = find_dropped_header_text(
            corrected_grid,
            cell_grid,
            " ".join(
                [correction.title, correction.caption, correction.units]
                + [f"{f.marker} {f.text}".strip() for f in correction.footnotes]
            ),
            page,
        )
    if correction is not None:
        return ExtractedTable(
            content_fingerprint=grid_fingerprint(cand.cells),
            title=correction.title,
            caption=correction.caption,
            markdown=correction.markdown,
            footnotes=correction.footnotes,
            footnote_refs=correction.footnote_refs,
            footnote_marks=locate_markers(
                correction.footnote_marks,
                correction.footnote_refs,
                correction.markdown,
                correction.footnotes,
                f"{correction.title} {correction.caption}",
            ),
            page=page,
            bbox=cand.bbox,
            flavor=cand.flavor,
            source=str(source),
            camelot_accuracy=cand.accuracy,
            llm_corrected=correction.llm_corrected,
            som_region=sub_region,
            units=correction.units,
            header_rows=correction.header_rows,
            cell_grid=cell_grid,
            corrected_grid=corrected_grid,
            merged_cells=merged,
            dropped_text=dropped_text,
        )
    return ExtractedTable(
        content_fingerprint=grid_fingerprint(cand.cells),
        markdown=grid_to_markdown(cand.cells),
        page=page,
        bbox=cand.bbox,
        flavor=cand.flavor,
        source=str(source),
        camelot_accuracy=cand.accuracy,
        som_region=sub_region,
        cell_grid=cell_grid,
        corrected_grid=corrected_grid,
    )


def render_page(source: Path, page: int, dpi: int, out_dir: Path) -> Tuple[Path, float, float]:
    """Render one PDF page to a PNG at `dpi`; return its path and page size in points."""
    doc = fitz.open(str(source))
    try:
        pg = doc[page - 1]
        width, height = pg.rect.width, pg.rect.height
        pix = pg.get_pixmap(dpi=dpi)
        path = out_dir / f"page-{page:04d}.png"
        pix.save(str(path))
    finally:
        doc.close()
    return path, width, height