"""
Camelot acquisition layer: get candidate cell grids out of a PDF.

This module knows nothing about LLMs or the pipeline graph. It renders
page images, runs both Camelot flavors in spawn-context subprocesses,
and converts each Camelot DataFrame into a serializable
`CamelotCandidate`.

Both Camelot flavors run in parallel:

- `lattice` — visible grid lines.
- `stream` — whitespace-based; more permissive.

There is deliberately no lattice-first / fall-back-to-stream gate: on
pages where tables are aligned by whitespace alone, lattice still
returns grids — empty shells — so "fall back only when lattice finds
nothing" let those shells through as false positives. `is_content_empty`
catches the shells instead.
"""

from __future__ import annotations

import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures import TimeoutError as FutureTimeoutError
from pathlib import Path
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple

from loguru import logger
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from pandas import DataFrame

Flavor = Literal["lattice", "stream"]

CAMELOT_FLAVOR_TIMEOUT_S = 60.0


def render_pages(source: Path, dpi: int, out_dir: Path) -> List[Path]:
    from pdf2image import convert_from_path

    images = convert_from_path(str(source), dpi=dpi, fmt="png", output_folder=str(out_dir))
    paths: List[Path] = []
    for i, img in enumerate(images, start=1):
        path = out_dir / f"page-{i:04d}.png"
        img.save(path, "PNG")
        paths.append(path)
    return paths


class CamelotCandidate(BaseModel):
    """Serializable per-Camelot-table output. Crosses the spawn-process
    boundary; only picklable primitives, no Camelot internals.

    `cells` is the raw grid (rows of cell strings) straight from Camelot's
    DataFrame; `markdown` is its rendering. The correspondence flow works
    on `cells` so it can assemble, audit and fill at the grid level and
    render markdown once at the end; the legacy classifier/unifier path
    consumes `markdown`.
    """

    candidate_id: str = Field(description="Stable id assigned by the parent; flavor-page-idx")
    flavor: Flavor
    page: int = Field(ge=1)
    bbox: Optional[Tuple[float, float, float, float]] = None
    accuracy: float = Field(default=0.0, ge=0.0, le=100.0)
    cells: List[List[str]] = Field(
        default_factory=list, description="Camelot's raw cell grid (rows of strings)"
    )
    cell_boxes: List[List[Optional[Tuple[float, float, float, float]]]] = Field(
        default_factory=list,
        description=(
            "Per-cell geometry aligned 1:1 with `cells`: same row/column shape, each entry the "
            "Camelot cell box (x1, y1, x2, y2) in PDF points with a bottom-left page origin, or "
            "None where Camelot exposed no box for that cell. Empty for candidates not built "
            "directly from a Camelot grid (e.g. assembled or recovered candidates)."
        ),
    )
    markdown: str = Field(description="Camelot's raw markdown for this table; rendering of `cells`")


def grid_to_markdown(cells: List[List[str]]) -> str:
    """Render a cell grid (rows of strings) to a markdown table. Row 0 is
    the header; rows are padded to the widest row so the column count is
    uniform. Single rendering point for the structured-grid flow.
    """
    if not cells:
        return ""
    width = max(len(r) for r in cells)

    def render(row: List[str]) -> str:
        padded = [str(c) for c in row] + [""] * (width - len(row))
        return "| " + " | ".join(padded) + " |"

    lines = [render(cells[0]), "| " + " | ".join("---" for _ in range(width)) + " |"]
    lines.extend(render(row) for row in cells[1:])
    return "\n".join(lines)


def column_letter(index: int) -> str:
    """0-based column index to a spreadsheet letter: 0->A, 25->Z, 26->AA."""
    out = ""
    index += 1
    while index:
        index, rem = divmod(index - 1, 26)
        out = chr(ord("A") + rem) + out
    return out


def grid_to_addressed_markdown(cells: List[List[str]]) -> str:
    """Render the grid with each non-empty cell's address printed inside it.

    Every cell with text carries an inline tag — `[B3] 1,637` means data
    column B (A=0), grid row 3 (1-based) — so the agent READS an address off
    the label sitting next to the text; it never counts rows, columns, or
    pipes. Blank cells stay blank. The tags are reference only — the agent's
    corrected output must never contain them.
    """
    if not cells:
        return ""
    width = max(len(r) for r in cells)

    def tag(r: int, c: int, text: str) -> str:
        return f"[{column_letter(c)}{r}] {text}" if str(text).strip() else ""

    def render_addressed(r: int, row: List[str]) -> str:
        padded = [str(c) for c in row] + [""] * (width - len(row))
        return "| " + " | ".join(tag(r, c, v) for c, v in enumerate(padded)) + " |"

    lines = [render_addressed(1, cells[0]), "| " + " | ".join("---" for _ in range(width)) + " |"]
    lines.extend(render_addressed(i, row) for i, row in enumerate(cells[1:], start=2))
    return "\n".join(lines)


def df_to_cells(df: DataFrame) -> List[List[str]]:
    """Camelot DataFrame -> raw cell grid (rows of strings)."""
    return [[str(c) for c in row] for row in df.values.tolist()]


def cells_to_boxes(
    raw_cells: Any, grid: List[List[str]]
) -> List[List[Optional[Tuple[float, float, float, float]]]]:
    """Per-cell boxes aligned to `grid`'s shape from Camelot's Cell objects.

    `raw_cells` is Camelot's `table.cells` — rows of `Cell` objects, each with
    x1/y1/x2/y2 in PDF points (bottom-left origin). `table.df` (hence `grid`) is
    built from that same grid, so index i,j lines up; still, we clamp to `grid`'s
    shape and fill None for any position Camelot did not cover, so `cell_boxes`
    is always exactly the same shape as `cells`.
    """
    rows = list(raw_cells) if raw_cells is not None else []
    boxes: List[List[Optional[Tuple[float, float, float, float]]]] = []
    for i, grid_row in enumerate(grid):
        raw_row = list(rows[i]) if i < len(rows) else []
        out_row: List[Optional[Tuple[float, float, float, float]]] = []
        for j in range(len(grid_row)):
            cell = raw_row[j] if j < len(raw_row) else None
            if cell is not None and all(hasattr(cell, a) for a in ("x1", "y1", "x2", "y2")):
                out_row.append((float(cell.x1), float(cell.y1), float(cell.x2), float(cell.y2)))
            else:
                out_row.append(None)
        boxes.append(out_row)
    return boxes


def df_to_markdown(df: DataFrame) -> str:
    return grid_to_markdown(df_to_cells(df))


def camelot_worker(source_str: str, flavor: Flavor) -> List[CamelotCandidate]:
    """Top-level entrypoint for the spawned process. Imports camelot
    fresh in the child so OpenCV initialization stays inside the worker
    and never crosses a fork boundary.
    """
    import camelot

    from quber.agents.completeness import page_words
    from quber.core.extractors.camelot.tighten import tighten_cell_boxes

    out: List[CamelotCandidate] = []
    tables = camelot.read_pdf(source_str, pages="all", flavor=flavor)  # pyright: ignore[reportPrivateImportUsage]
    # Word rectangles per page, read once and shared by every table on it.
    words_cache: dict[int, Tuple[float, List[Any]]] = {}
    for idx, table in enumerate(tables):
        raw_page = getattr(table, "page", 1)
        page = int(raw_page) if raw_page is not None else 1
        # camelot exposes the table box only as the private `_bbox`; read it
        # via getattr so the access is not flagged as private use.
        raw_bbox = getattr(table, "_bbox", None)
        bbox = tuple(raw_bbox) if raw_bbox else None
        report = getattr(table, "parsing_report", {}) or {}
        accuracy = float(report.get("accuracy", 0.0) or 0.0)
        cells = df_to_cells(table.df)
        cell_boxes = cells_to_boxes(getattr(table, "cells", None), cells)
        if page not in words_cache:
            _, page_h, words = page_words(Path(source_str), page)
            words_cache[page] = (page_h, words)
        page_h, words = words_cache[page]
        cell_boxes = tighten_cell_boxes(cells, cell_boxes, words, page_h)
        out.append(
            CamelotCandidate(
                candidate_id=f"{flavor}-p{page}-{idx}",
                flavor=flavor,
                page=page,
                bbox=bbox,
                accuracy=accuracy,
                cells=cells,
                cell_boxes=cell_boxes,
                markdown=grid_to_markdown(cells),
            )
        )
    return out


def run_camelot_flavors_parallel(
    source: Path,
    timeout_s: float = CAMELOT_FLAVOR_TIMEOUT_S,
) -> List[CamelotCandidate]:
    """Run lattice + stream concurrently in spawn-context subprocesses.

    A flavor that times out or raises is logged ERROR and contributes
    zero candidates; the other flavor's output is still returned. If
    BOTH fail, returns an empty list — the caller decides what to do.
    """
    ctx = mp.get_context("spawn")
    candidates: List[CamelotCandidate] = []
    flavors: List[Flavor] = ["lattice", "stream"]

    with ProcessPoolExecutor(max_workers=2, mp_context=ctx) as pool:
        futures = {flavor: pool.submit(camelot_worker, str(source), flavor) for flavor in flavors}
        for flavor, future in futures.items():
            try:
                candidates.extend(future.result(timeout=timeout_s))
            except FutureTimeoutError:
                logger.error(
                    "camelot {flavor} timed out after {timeout}s on {source}",
                    flavor=flavor,
                    timeout=timeout_s,
                    source=source.name,
                )
                future.cancel()
            except Exception as exc:
                logger.error(
                    "camelot {flavor} raised on {source}: {exc}",
                    flavor=flavor,
                    source=source.name,
                    exc=exc,
                )

    if not candidates:
        logger.error("camelot produced 0 candidates from BOTH flavors on {}", source.name)
    return candidates


def is_content_empty(markdown: str) -> bool:
    """A Camelot output is content-empty when every non-delimiter cell
    is whitespace. Lattice's false-positive shells on whitespace-aligned
    PDFs look like a 1-row grid of empty pipes; this catches them
    before they reach the classifier or correction stage.
    """
    if not markdown.strip():
        return True
    for line in markdown.splitlines():
        cells = [c.strip() for c in line.strip().strip("|").split("|")]
        if any(c and not set(c) <= {"-", " "} for c in cells):
            return False
    return True