"""Read a scanned page's tables back as grids.

A scan returns a table as a table. The chunk carries an HTML grid with an id on
every cell, and the response's grounding map carries, against each of those ids,
the cell's box and its row, column and spans. Putting the two together gives a
dense grid — one entry per row and column — with a box on every cell, which is
the shape the rest of the table pipeline already works in.

A table is not always its own chunk. Where a page prints a chart and a table as
one composite, the scan returns one figure, and writes the table inside that
figure's text as markdown between a pair of markers it emits itself. The values
are complete and correctly separated there — on one page the scan read a bond
issuance table's four credit-rating rows exactly, while the parse ran four
ratings together into a single cell and left the next three blank.

So the markers are read and the markdown between them is taken as a table.
Nothing infers where the table is in the text; the scan labelled it.

What such a table does not carry is geometry. The composite is grounded as one
box covering the whole figure, with nothing per cell, so its cells arrive with no
boxes. That costs nothing downstream: the structure correction reads the page
image rather than the boxes, and a table with no reading to check against already
puts every cell to the agent.

Nothing here interprets the content. Spans are honoured as geometry: a cell that
spans two columns is written once at its own position and the positions it covers
are left blank, the same way a grid extracted from a text layer arrives. What the
header block is, which cells merged, and what the title says are read later, off
the table image, by the steps that read every other table.
"""

from __future__ import annotations

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

from quber.core.extractors.set_of_mark.merge_grounding import markdown_rows
from quber.core.figures.models import Box, ScannedTable

#: The grounding entry type carrying one table cell's box and grid position.
CELL_TYPE = "tableCell"

#: The grounding entry type carrying the box around a whole table. It is drawn
#: tighter than the chunk's own box, which reaches out to the surrounding text.
TABLE_TYPE = "table"

_TABLE_ID_RE = re.compile(r"<table[^>]*\bid=[\"']([^\"']+)[\"']")

#: The markers the scan writes around a markdown table it printed inside a
#: figure's text, when the page prints a chart and a table as one composite.
_FIGURE_TABLE_RE = re.compile(r"<::table::>(.*?)<::/table::>", re.DOTALL)


def scanned_tables(response: Dict[str, Any], page: int, job_id: Optional[str]) -> List[ScannedTable]:
    """Every table in one page's raw response, as a dense grid with per-cell boxes.

    `page` is the 1-based source page the submission came from. The response
    numbers its own single page from zero, so the source page is carried in
    rather than read back out of it.
    """
    grounding = response.get("grounding") or {}
    cells_by_chunk = _cells_by_chunk(grounding)

    tables: List[ScannedTable] = []
    for chunk in response.get("chunks") or []:
        if chunk.get("type") != "table":
            continue
        chunk_id = chunk.get("id")
        markdown = chunk.get("markdown") or ""
        text_by_id = _cell_text(markdown)
        cells, boxes = _dense_grid(cells_by_chunk.get(chunk_id, []), text_by_id)
        tables.append(
            ScannedTable(
                page=page,
                cells=cells,
                cell_boxes=boxes,
                box=_table_box(markdown, grounding) or (chunk.get("grounding") or {}).get("box"),
                chunk_id=chunk_id,
                job_id=job_id,
            )
        )

    tables.extend(_composite_tables(response, page, job_id))
    return tables


def _composite_tables(response: Dict[str, Any], page: int, job_id: Optional[str]) -> List[ScannedTable]:
    """The tables the scan printed inside a figure's text rather than as chunks.

    The figure's own box stands as the table's box. It is the only geometry the
    scan grounded for the composite, and it is drawn around the chart and the
    table together, so the region is generous rather than wrong.
    """
    out: List[ScannedTable] = []
    for chunk in response.get("chunks") or []:
        if chunk.get("type") != "figure":
            continue
        for markdown in _FIGURE_TABLE_RE.findall(chunk.get("markdown") or ""):
            rows = markdown_rows(markdown)
            if not rows:
                continue
            width = max(len(r) for r in rows)
            cells = [r + [""] * (width - len(r)) for r in rows]
            out.append(
                ScannedTable(
                    page=page,
                    cells=cells,
                    cell_boxes=[[None] * width for _ in cells],
                    box=(chunk.get("grounding") or {}).get("box"),
                    chunk_id=chunk.get("id"),
                    job_id=job_id,
                )
            )
    return out


def _cells_by_chunk(grounding: Dict[str, Any]) -> Dict[str, List[Tuple[str, Dict[str, Any]]]]:
    """The grounding map's cell entries, grouped by the table chunk each belongs to."""
    by_chunk: Dict[str, List[Tuple[str, Dict[str, Any]]]] = {}
    for cell_id, entry in grounding.items():
        if entry.get("type") != CELL_TYPE:
            continue
        position = entry.get("position") or {}
        chunk_id = position.get("chunk_id")
        if chunk_id is None:
            continue
        by_chunk.setdefault(chunk_id, []).append((cell_id, entry))
    return by_chunk


def _table_box(markdown: str, grounding: Dict[str, Any]) -> Optional[Box]:
    """The box around the table itself, from the id its opening tag carries."""
    match = _TABLE_ID_RE.search(markdown)
    if match is None:
        return None
    entry = grounding.get(match.group(1)) or {}
    return entry.get("box") if entry.get("type") == TABLE_TYPE else None


def _dense_grid(
    cells: List[Tuple[str, Dict[str, Any]]],
    text_by_id: Dict[str, str],
) -> Tuple[List[List[str]], List[List[Optional[Box]]]]:
    """One row per row and one column per column, with a box beside every cell.

    A cell that spans is written at its own row and column; the positions it
    covers stay blank and carry no box of their own, so the grid never states the
    same value twice.
    """
    placed: List[Tuple[int, int, str, Optional[Box]]] = []
    rows = 0
    columns = 0
    for cell_id, entry in cells:
        position = entry.get("position") or {}
        row = int(position.get("row", 0))
        col = int(position.get("col", 0))
        rows = max(rows, row + int(position.get("rowspan", 1) or 1))
        columns = max(columns, col + int(position.get("colspan", 1) or 1))
        placed.append((row, col, text_by_id.get(cell_id, ""), entry.get("box")))

    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
    return grid, boxes


class _CellText(HTMLParser):
    """Collects the text of every `td`/`th` carrying an id, keyed by that id.

    Markup inside a cell is dropped and its text kept, so a value typeset with an
    emphasis or a line break reads as the one string the page prints.
    """

    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.text: Dict[str, str] = {}
        self._open: List[Optional[str]] = []
        self._parts: List[str] = []

    def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None:
        if tag in ("td", "th"):
            self._flush()
            self._open.append(dict(attrs).get("id"))
        elif tag == "br" and self._open:
            self._parts.append(" ")

    def handle_endtag(self, tag: str) -> None:
        if tag in ("td", "th"):
            self._flush()

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

    def _flush(self) -> None:
        if not self._open:
            return
        cell_id = self._open.pop()
        if cell_id is not None:
            self.text[cell_id] = " ".join("".join(self._parts).split())
        self._parts = []


def _cell_text(markdown: str) -> Dict[str, str]:
    """Each cell's printed text, keyed by the id its tag carries."""
    parser = _CellText()
    parser.feed(markdown)
    parser.close()
    return parser.text