"""Layout metadata: reading order, heading binding, image context."""

from __future__ import annotations

from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.tables import build_table_embedding_text, extract_table_keywords, guess_table_title
from rag_cut.splitters.heading_splitter import is_heading_block, normalize_heading_block


def assign_order_index(blocks: list[Block], start: int = 0) -> list[Block]:
    """Assign a global document-order index to every block."""
    result: list[Block] = []
    for idx, block in enumerate(blocks):
        meta = dict(block.meta)
        meta["order_index"] = start + idx
        result.append(block.model_copy(update={"meta": meta}))
    return result


def sort_blocks_reading_order(blocks: list[Block]) -> list[Block]:
    """Sort by page then top-to-bottom / left-to-right; stable for missing bboxes."""

    def sort_key(item: tuple[int, Block]) -> tuple[int, float, float, int]:
        idx, block = item
        page = block.meta.get("page")
        try:
            page_key = int(page) if page is not None else 10**9
        except (TypeError, ValueError):
            page_key = 10**9
        bbox = block.meta.get("bbox")
        if isinstance(bbox, (list, tuple)) and len(bbox) >= 4:
            try:
                return (page_key, float(bbox[1]), float(bbox[0]), idx)
            except (TypeError, ValueError):
                pass
        return (page_key, float(idx), 0.0, idx)

    return [block for _, block in sorted(enumerate(blocks), key=sort_key)]


def bind_heading_context(blocks: list[Block]) -> list[Block]:
    """
    Propagate heading hierarchy and bind images to nearest heading and adjacent text.
    Chunk content should follow: heading → body → image → OCR/caption → subsequent body.
    """
    heading_stack: list[tuple[int, str]] = []
    result: list[Block] = []

    for i, block in enumerate(blocks):
        meta = dict(block.meta)
        candidate = normalize_heading_block(block)

        if is_heading_block(candidate):
            level = candidate.level or 1
            while heading_stack and heading_stack[-1][0] >= level:
                heading_stack.pop()
            heading_stack.append((level, candidate.text))
            meta["section_boundary"] = True
            if block.type != BlockType.HEADING or candidate is not block:
                block = candidate.model_copy(update={"meta": {**dict(candidate.meta), **meta}})
                meta = dict(block.meta)
        else:
            meta["section_boundary"] = False
            # Demote overlong HEADING blobs (title+body merge) back to paragraph.
            if block.type == BlockType.HEADING and candidate.type != BlockType.HEADING:
                block = candidate.model_copy(update={"meta": {**dict(candidate.meta), **meta}})
                meta = dict(block.meta)

        if heading_stack:
            meta["parent_heading"] = heading_stack[-1][1]
            meta["heading_path"] = [text for _, text in heading_stack]
            meta["nearest_heading"] = heading_stack[-1][1]
            if (block.level or 1) <= 2 and block.type == BlockType.HEADING:
                meta["chapter"] = block.text
            elif "chapter" not in meta and len(heading_stack) >= 1:
                # Keep chapter as nearest level-1/2 ancestor
                for lvl, text in reversed(heading_stack):
                    if lvl <= 2:
                        meta["chapter"] = text
                        break

        if block.type == BlockType.IMAGE:
            _bind_image_context(blocks, i, meta)
        elif block.type == BlockType.TABLE:
            _bind_table_context(blocks, i, meta)

        result.append(block.model_copy(update={"meta": meta}))

    return result


def _heading_text(block: Block) -> str | None:
    candidate = normalize_heading_block(block)
    if is_heading_block(candidate):
        return (candidate.text or "").strip() or None
    return None


def _spatial_heading_above(
    blocks: list[Block],
    index: int,
    page: object,
    bbox: list[float] | None,
) -> str | None:
    """Pick same-page heading whose bottom edge is nearest above the image top."""
    if page is None or not bbox or len(bbox) < 4:
        return None
    try:
        page_key = int(page)
        img_y0 = float(bbox[1])
    except (TypeError, ValueError):
        return None

    best_text: str | None = None
    best_dist = float("inf")
    for j, block in enumerate(blocks):
        if j == index:
            continue
        text = _heading_text(block)
        if not text:
            continue
        try:
            if int(block.meta.get("page")) != page_key:
                continue
        except (TypeError, ValueError):
            continue
        hb = block.meta.get("bbox")
        if not isinstance(hb, (list, tuple)) or len(hb) < 4:
            continue
        try:
            heading_y1 = float(hb[3])
        except (TypeError, ValueError):
            continue
        if heading_y1 > img_y0 + 2:
            continue
        dist = img_y0 - heading_y1
        if dist < best_dist:
            best_dist = dist
            best_text = text
    return best_text


def _list_heading_above(blocks: list[Block], index: int) -> str | None:
    for j in range(index - 1, -1, -1):
        text = _heading_text(blocks[j])
        if text:
            return text
    return None


def _bind_image_context(blocks: list[Block], index: int, meta: dict) -> None:
    """Bind image to nearest heading above and adjacent body text on the same page."""
    page = meta.get("page")
    bbox = meta.get("bbox") if isinstance(meta.get("bbox"), list) else None

    bound = _spatial_heading_above(blocks, index, page, bbox) or _list_heading_above(blocks, index)
    if bound:
        meta["bound_heading"] = bound

    for j in range(index - 1, -1, -1):
        prev = blocks[j]
        if _heading_text(prev):
            break
        if prev.type == BlockType.PARAGRAPH and prev.meta.get("page") == page:
            meta["preceding_text"] = (prev.text or "")[:300]
            break

    if not meta.get("bound_heading") and meta.get("nearest_heading"):
        meta["bound_heading"] = meta["nearest_heading"]

    for j in range(index + 1, len(blocks)):
        nxt = blocks[j]
        if nxt.type == BlockType.IMAGE:
            break
        if nxt.type == BlockType.PARAGRAPH and nxt.meta.get("page") == page:
            meta["following_text"] = (nxt.text or "")[:300]
            break
        if _heading_text(nxt):
            break


def _bind_table_context(blocks: list[Block], index: int, meta: dict) -> None:
    """Bind table title, surrounding text and retrieval fields."""
    page = meta.get("page")
    block = blocks[index]

    caption = (meta.get("caption") or "").strip()
    if caption and len(caption) <= 160 and not meta.get("table_title"):
        meta["table_title"] = caption

    for j in range(index - 1, -1, -1):
        prev = blocks[j]
        if prev.type == BlockType.TABLE:
            break
        heading = _heading_text(prev)
        if heading:
            if not meta.get("table_title"):
                meta["table_title"] = heading
            meta.setdefault("bound_heading", heading)
            break
        if prev.type == BlockType.PARAGRAPH and prev.meta.get("page") == page:
            title = guess_table_title(prev.text or "")
            if title:
                meta["table_title"] = title
            meta["preceding_text"] = (prev.text or "")[:400]
            break

    if not meta.get("table_title") and meta.get("nearest_heading"):
        meta.setdefault("table_title", meta["nearest_heading"])

    for j in range(index + 1, len(blocks)):
        nxt = blocks[j]
        if nxt.type == BlockType.TABLE:
            break
        if nxt.type == BlockType.PARAGRAPH and nxt.meta.get("page") == page:
            text = (nxt.text or "").strip()
            if text and len(text) <= 300:
                meta["following_text"] = text
                if not meta.get("footnotes") and any(k in text for k in ("注", "备注", "说明", "Note")):
                    meta["table_description"] = text
            break
        if _heading_text(nxt):
            break

    if not meta.get("keywords"):
        meta["keywords"] = extract_table_keywords(
            meta.get("table_title") or "",
            block.markdown or "",
            meta.get("table_description") or meta.get("footnotes") or "",
            block.ocr_text or "",
        )

    meta["embedding_text"] = build_table_embedding_text(
        chapter=meta.get("chapter") or meta.get("nearest_heading") or "",
        table_title=meta.get("table_title") or "",
        markdown=block.markdown or "",
        description=meta.get("table_description") or meta.get("preceding_text") or "",
        footnotes=meta.get("footnotes") or meta.get("following_text") or "",
        keywords=meta.get("keywords") or [],
        ocr_text=block.ocr_text or "",
    )


def enrich_layout_metadata(blocks: list[Block]) -> list[Block]:
    """Full post-parse enrichment: spatial order + order index + heading/image context."""
    blocks = sort_blocks_reading_order(blocks)
    blocks = assign_order_index(blocks)
    return bind_heading_context(blocks)