"""Table-aware chunk grouping helpers."""

from __future__ import annotations

from rag_cut.models import Block, BlockType


def group_tables_with_context(blocks: list[Block]) -> list[list[Block]]:
    """
    Group each table with its contextual title paragraph and footnote.
    Keeps: heading/intro → table → footnote as one atomic group when adjacent.
    """
    if not blocks:
        return []

    groups: list[list[Block]] = []
    current: list[Block] = []
    i = 0

    while i < len(blocks):
        block = blocks[i]
        if block.type != BlockType.TABLE:
            current.append(block)
            i += 1
            continue

        group: list[Block] = []
        if current:
            tail = current[-1]
            if tail.type in {BlockType.HEADING, BlockType.PARAGRAPH}:
                group.append(tail)
                current = current[:-1]
            if current:
                groups.append(current)
                current = []

        group.append(block)
        j = i + 1
        while j < len(blocks):
            nxt = blocks[j]
            if nxt.type == BlockType.PARAGRAPH and len((nxt.text or "")) <= 320:
                if any(k in (nxt.text or "") for k in ("注", "备注", "说明", "Note", "※")):
                    group.append(nxt)
                    j += 1
                    break
            break

        groups.append(group)
        i = j

    if current:
        groups.append(current)
    return groups