"""Turn one page's dpt-3 response into records, and project them for the graft.

A dpt-3 response has two parts: one long markdown string holding all the text
the model read off the page, and a tree describing what sits where. Each tree
node names its kind, gives its rectangle on the page, and points at the stretch
of the long string that belongs to it — the node itself stores no text, so
getting a node's text means cutting its stretch out of the string.

`digest_page` walks the tree and produces the track's records: a figure node
becomes a `DigestedFigure` with its kind, description and value tables kept
separate; a table node and its cell nodes become a `DigestedTable` with a box
on every cell; everything else on the page is kept as context. A response
without the tree raises `UnrecognizedResponse` naming the stored file and the
model, so a billed page can never pass as blank because the code did not
recognize the format. A page node reporting anything but success raises too,
with the response's own reason.

`project_scan` turns a digest into the `PageScan` every downstream consumer
already reads — the graft, the correction sweep, the value reconciliation. A
figure is projected as its description followed by its value tables rendered
as markdown; a table is projected grid-for-grid. The rest of the digest stays
in the track models.

The figure markup is the response's own labelling — the model wraps each
figure's content in tags it emits itself — so reading it back is reading the
format, not inferring anything about the page.
"""

from __future__ import annotations

import re
from html.parser import HTMLParser
from typing import Any, Dict, List, Optional, Tuple

from quber.core.figures.dpt3.models import (
    DigestedFigure,
    DigestedTable,
    DigestedText,
    FigureValueCell,
    FigureValueTable,
    PageDigest,
)
from quber.core.figures.models import Box, ChartContext, FigureRecord, PageScan, PageStatus, ScannedTable


class UnrecognizedResponse(ValueError):
    """A response in neither the dpt-2 nor the dpt-3 format."""


#: The opening tag the response wraps a figure's content in, carrying its kind.
_FIGURE_OPEN_RE = re.compile(r'<figure\b[^>]*\btype="([^"]*)"[^>]*>')

#: One description block inside a figure's content.
_DESCRIPTION_RE = re.compile(r"<description>(.*?)</description>", re.DOTALL)

#: One value table inside a figure's content.
_VALUE_TABLE_RE = re.compile(r"<table>(.*?)</table>", re.DOTALL)


def response_generation(response: Dict[str, Any]) -> Optional[str]:
    """Which format a raw response holds: `dpt2`, `dpt3`, or None for neither.

    The format is visible in the response itself: a dpt-2 response carries a
    flat `chunks` list, a dpt-3 response carries the `structure` tree.
    """
    if isinstance(response.get("chunks"), list):
        return "dpt2"
    if isinstance(response.get("structure"), dict):
        return "dpt3"
    return None


def digest_page(
    response: Dict[str, Any],
    page: int,
    model: str,
    artifact: Optional[str] = None,
) -> PageDigest:
    """Digest one page's raw dpt-3 response into the track's records.

    `page` is the 1-based source page the submission came from; the response
    numbers its own single page from one, so the source page is carried in
    rather than read back out of it. `artifact` names the stored response file
    for error messages.
    """
    structure = response.get("structure")
    if not isinstance(structure, dict):
        raise UnrecognizedResponse(
            f"{artifact or 'the response'} (model {model}) holds neither a dpt-2 chunks list "
            "nor a dpt-3 structure tree"
        )
    markdown = response.get("markdown") or ""
    meta = response.get("metadata") or {}

    digest = PageDigest(
        page=page,
        job_id=meta.get("job_id"),
        model=model,
        version=meta.get("model_version") or meta.get("version"),
        credits=(meta.get("billing") or {}).get("total_credits", meta.get("credit_usage")),
    )
    for page_node in structure.get("children") or []:
        if page_node.get("type") != "page":
            continue
        status = page_node.get("status")
        if status not in (None, "ok"):
            raise ValueError(
                f"{artifact or 'the response'} (model {model}) reports page status {status!r}"
                f" ({page_node.get('reason')}); the page was billed and cannot pass as blank"
            )
        for node in page_node.get("children") or []:
            kind = node.get("type")
            if kind == "figure":
                digest.figures.append(_figure(node, markdown, page))
            elif kind == "table":
                digest.tables.append(_table(node, markdown, page))
            else:
                digest.context.append(
                    DigestedText(
                        page=page,
                        kind=kind or "",
                        text=_slice(node, markdown),
                        box=_box(node),
                        node_id=node.get("id"),
                    )
                )
    return digest


def project_scan(
    digest: PageDigest,
    picture_classes: List[str],
    response_artifact: Optional[str] = None,
    reused: bool = False,
) -> PageScan:
    """The digest as the `PageScan` every downstream consumer already reads.

    The figure records carry text, id and rectangle — the graft's whole
    contract. The digest keeps everything else.
    """
    figures = [
        FigureRecord(
            page=digest.page,
            text=figure_text(figure),
            box=figure.box,
            chunk_id=figure.node_id,
            job_id=digest.job_id,
        )
        for figure in digest.figures
    ]
    tables = [
        ScannedTable(
            page=digest.page,
            cells=table.cells,
            cell_boxes=table.cell_boxes,
            box=table.box,
            chunk_id=table.node_id,
            job_id=digest.job_id,
        )
        for table in digest.tables
    ]
    context = [
        ChartContext(page=digest.page, kind=item.kind, text=item.text, box=item.box, chunk_id=item.node_id)
        for item in digest.context
        if item.text.strip()
    ]
    status: PageStatus = "figures" if figures else "tables" if tables else "empty"
    return PageScan(
        page=digest.page,
        status=status,
        picture_classes=picture_classes,
        job_id=digest.job_id,
        model=digest.model,
        version=digest.version,
        credits=digest.credits,
        reused=reused,
        response_artifact=response_artifact,
        figures=figures,
        tables=tables,
        context=context,
    )


def figure_text(figure: DigestedFigure) -> str:
    """A figure's readable content: its description, then its value tables.

    This is what the graft puts on the picture and what the correction agent
    and the value reconciliation read, so the values are rendered as markdown
    rows rather than left in the response's markup.
    """
    parts: List[str] = []
    if figure.description.strip():
        parts.append(figure.description.strip())
    for table in figure.values:
        rendered = _markdown_rows(table)
        if rendered:
            parts.append(rendered)
    return "\n\n".join(parts)


def _markdown_rows(table: FigureValueTable) -> str:
    """A value table as pipe-delimited rows, one line per row."""
    lines = []
    for row in table.rows:
        cells = [" ".join(cell.text.split()).replace("|", "/") for cell in row]
        lines.append("| " + " | ".join(cells) + " |")
        if len(lines) == 1 and len(table.rows) > 1:
            lines.append("|" + " --- |" * len(row))
    return "\n".join(lines)


def _figure(node: Dict[str, Any], markdown: str, page: int) -> DigestedFigure:
    """One figure node: kind, descriptions and value tables out of its stretch."""
    content = _slice(node, markdown)
    kind_match = _FIGURE_OPEN_RE.search(content)
    descriptions = [d.strip() for d in _DESCRIPTION_RE.findall(content) if d.strip()]
    tables = [_value_table(html) for html in _VALUE_TABLE_RE.findall(content)]
    return DigestedFigure(
        page=page,
        kind=(kind_match.group(1).lower() if kind_match else ""),
        description="\n\n".join(descriptions),
        box=_box(node),
        values=[t for t in tables if t.rows],
        node_id=node.get("id"),
    )


def _table(node: Dict[str, Any], markdown: str, page: int) -> DigestedTable:
    """One table node as a dense grid, cells and boxes read off its cell nodes.

    A cell that spans is written at its own row and column; the positions it
    covers stay blank and carry no box, so the grid never states a value twice.
    """
    placed: List[Tuple[int, int, str, Optional[Box]]] = []
    rows = 0
    columns = 0
    for cell in node.get("children") or []:
        if cell.get("type") != "table_cell":
            continue
        row = int(cell.get("row") or 0)
        col = int(cell.get("col") or 0)
        rows = max(rows, row + int(cell.get("rowspan") or 1))
        columns = max(columns, col + int(cell.get("colspan") or 1))
        placed.append((row, col, " ".join(_slice(cell, markdown).split()), _box(cell)))

    grid = [["" for _ in range(columns)] for _ in range(rows)]
    boxes: List[List[Optional[Box]]] = [[None for _ in range(columns)] for _ in range(rows)]
    for row, col, text, box in placed:
        if 0 <= row < rows and 0 <= col < columns:
            grid[row][col] = text
            boxes[row][col] = box

    if not placed:
        # A table node without cell nodes still states its grid in its stretch.
        value_table = _value_table(_slice(node, markdown))
        grid = [[cell.text for cell in row] for row in value_table.rows]
        boxes = [[None for _ in row] for row in grid]

    return DigestedTable(page=page, cells=grid, cell_boxes=boxes, box=_box(node), node_id=node.get("id"))


def _value_table(html: str) -> FigureValueTable:
    """One value table's rows, read out of the response's own table markup."""
    parser = _GridText()
    parser.feed(html)
    parser.close()
    return FigureValueTable(rows=[[FigureValueCell(text=cell) for cell in row] for row in parser.rows if row])


class _GridText(HTMLParser):
    """Collects `td`/`th` text row by row. Markup inside a cell is dropped and
    its text kept, so a value typeset with a line break reads as one string."""

    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.rows: List[List[str]] = []
        self._row: Optional[List[str]] = None
        self._parts: Optional[List[str]] = None

    def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None:
        if tag == "tr":
            self._flush_row()
            self._row = []
        elif tag in ("td", "th"):
            self._flush_cell()
            self._parts = []
        elif tag == "br" and self._parts is not None:
            self._parts.append(" ")

    def handle_endtag(self, tag: str) -> None:
        if tag in ("td", "th"):
            self._flush_cell()
        elif tag == "tr":
            self._flush_row()

    def handle_data(self, data: str) -> None:
        if self._parts is not None:
            self._parts.append(data)

    def close(self) -> None:
        super().close()
        self._flush_row()

    def _flush_cell(self) -> None:
        if self._parts is None:
            return
        if self._row is None:
            self._row = []
        self._row.append(" ".join("".join(self._parts).split()))
        self._parts = None

    def _flush_row(self) -> None:
        self._flush_cell()
        if self._row is not None:
            self.rows.append(self._row)
            self._row = None


def _slice(node: Dict[str, Any], markdown: str) -> str:
    """A node's text: the stretch of the response's one long string it points at."""
    grounding = node.get("grounding") or {}
    rng = grounding.get("range") or {}
    start, end = rng.get("start"), rng.get("end")
    if start is None or end is None:
        return ""
    return markdown[start:end]


def _box(node: Dict[str, Any]) -> Optional[Box]:
    """A node's rectangle in the workflow's own box keys, or None without one."""
    grounding = node.get("grounding") or {}
    box = grounding.get("box") or {}
    try:
        return {
            "left": box["xmin"],
            "top": box["ymin"],
            "right": box["xmax"],
            "bottom": box["ymax"],
        }
    except KeyError:
        return None