"""Group the parse's floating text by the scan's block partition.

A page like a stat-panel overview prints its content as short label and value
lines. The parse captures every line, but as separate records — `$78B` is one
record, `Investor capital` another, joined by nothing — and anything consuming
the records one at a time hands a reader a label without its number. The scan's
response partitions the same page into blocks: every line belongs to exactly
one node, and the lines of one panel come back inside one node.

This module carries that partition onto the parse's own elements. Each plain
text item on a scanned page is assigned to the block whose rectangle holds its
center, and the items of a block are re-parented under one group node in the
document, in their reading order, at the first member's position in the flow.
The items themselves are untouched: their text, boxes and references stay the
parse's own, and nothing is transcribed from the scan. The scan contributes one
fact the parse does not state — which lines sit together.

The blocks come from the response's text nodes and from its table nodes. The
scan reads a panel collage as a table as readily as it reads one as text, and
a printed panel is the same panel either way. A table node contributes each of
its cells as a block, so a value and its label bind at the cell, not across
the whole collage; the node's own rectangle stands behind the cells for text
that sits between them, and the smallest containing block always wins.

The same division of authority the graft enforces for figures: the scan looked
at the page, so it decides the partition; the parse's elements supply the
content.

Only plain text items parented to the document body take part. An item under a
picture is the graft's business, a heading or a footnote or a caption has a
role the partition must not swallow, and an item already inside a group is
left alone, whether the parse made the group (a list, an inline group) or an
earlier run did. That is why a rerun over an enriched parse changes nothing.
A block holding fewer than two items gains nothing from a group and gets none.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Sequence, Tuple

from docling_core.types.doc.document import DoclingDocument, TextItem
from docling_core.types.doc.items.node import NodeItem
from docling_core.types.doc.labels import DocItemLabel, GroupLabel
from loguru import logger

from quber.core.figures.dpt3.models import PageDigest
from quber.core.figures.geometry import NormBox, norm_box, prov_box
from quber.core.figures.models import PageScan

#: The name every group written by this module carries, so a consumer chunking
#: the document (`quber.playground.ingest_fusion`) can treat a scan block as
#: one unit. A rerun does not read it: it skips every item inside any group.
SCAN_BLOCK = "scan-text-block"


def group_scanned_text(
    document: DoclingDocument,
    scans: Sequence[PageScan],
    digests: Dict[int, PageDigest],
    page_dims: Optional[Dict[int, Tuple[float, float]]] = None,
) -> int:
    """Group each scanned page's plain text by the scan's block partition.

    Mutates `document`: the items of each block move under one group node at
    the first member's reading-order position. Returns how many groups were
    written.
    """
    dims = dict(page_dims or {})
    for page_no, page in document.pages.items():
        if page_no not in dims and page.size is not None:
            dims[page_no] = (page.size.width, page.size.height)

    written = 0
    for scan in scans:
        digest = digests.get(scan.page)
        if digest is None:
            continue
        blocks = [
            box
            for box in (norm_box(item.box) for item in digest.context if item.kind == "text")
            if box is not None
        ]
        for table in digest.tables:
            blocks.extend(box for row in table.cell_boxes for box in map(norm_box, row) if box is not None)
            table_box = norm_box(table.box)
            if table_box is not None:
                blocks.append(table_box)
        if not blocks:
            continue
        width, height = dims.get(scan.page, (612.0, 792.0))
        members: Dict[int, List[TextItem]] = {}
        for item in _floating_text(document, scan.page):
            box = prov_box(item, width, height)
            block = _owning_block(box, blocks)
            if block is not None:
                members.setdefault(block, []).append(item)

        for block_index in sorted(members):
            items = members[block_index]
            if len(items) < 2:
                continue
            group = document.insert_group(
                sibling=items[0], label=GroupLabel.UNSPECIFIED, name=SCAN_BLOCK, after=False
            )
            for item in items:
                _reparent(document, item, group)
            written += 1
            logger.info(
                "Text blocks: page {} grouped {} item(s) ({})",
                scan.page,
                len(items),
                " / ".join((item.text or "")[:24] for item in items[:4]),
            )
    return written


def _floating_text(document: DoclingDocument, page: int) -> List[TextItem]:
    """The page's plain text items parented to the body, in document order.

    Everything else keeps its place: an item under a picture belongs to the
    graft, a labelled item has a role, and an item already inside a group,
    the parse's own or an earlier run's, keeps its group.
    """
    body = document.body
    found: List[TextItem] = []
    for item in document.texts:
        if item.label != DocItemLabel.TEXT or not (item.text or "").strip():
            continue
        if not item.prov or item.prov[0].page_no != page:
            continue
        parent = item.parent.resolve(document) if item.parent else None
        if parent is not body:
            continue
        found.append(item)
    return found


def _owning_block(box: Optional[NormBox], blocks: List[NormBox]) -> Optional[int]:
    """The block whose rectangle holds the item's center; the smallest when
    blocks overlap; None when no block claims it."""
    if box is None:
        return None
    cx, cy = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2
    best: Optional[Tuple[float, int]] = None
    for i, (x1, y1, x2, y2) in enumerate(blocks):
        if x1 <= cx <= x2 and y1 <= cy <= y2:
            area = (x2 - x1) * (y2 - y1)
            if best is None or area < best[0]:
                best = (area, i)
    return best[1] if best is not None else None


def _reparent(document: DoclingDocument, item: TextItem, group: NodeItem) -> None:
    """Move one item under the group, keeping its identity and its order.

    Reference surgery rather than delete-and-re-add: the item keeps its
    self_ref, so nothing else in the document renumbers.
    """
    old_parent = item.parent.resolve(document) if item.parent else None
    if old_parent is not None:
        old_parent.children = [ref for ref in old_parent.children if ref.cref != item.self_ref]
    group.children.append(item.get_ref())
    item.parent = group.get_ref()