"""Document heading review — demote page decoration the parser called a heading.

The parser labels certain printed lines as section headings, and downstream
consumers treat a heading as governing every chunk beneath it. When the label
is wrong — a banner the page layout repeats at the top of its pages, or a
table column label that leaked into the heading stream — that false context
spreads across a large share of the document's text.

Demotion is deliberately hard, because wrongly demoting a real heading strips
section context from everything under it. Two witnesses must agree:

1. Geometry nominates: the same text appears as a heading on
   ``NOMINATION_PAGES`` or more pages, computed from the parse.
2. The heading-review agent confirms, judging only the nominated texts.

A heading that appears on fewer pages is never put to the agent, and an
unavailable review demotes nothing.

One demotion follows from another instead of from a verdict. The parser can
split a running header in two and label its first half a heading. Once the
agent confirms a text as a running header, any heading occurrence whose text
opens that header and is printed within ``SPLIT_HEADER_TOLERANCE`` of the
header's position band is demoted too, whatever its page count.

Every demotion is returned as a flag record, the run's record that the relabel
happened and why. A verdict demotion names the text, its pages, the verdict and
reason, and how many text items it governed. A split-header demotion names the
page and where the line is printed against the header's band.
"""

from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING, Dict, List, Set

from docling_core.types.doc.labels import DocItemLabel
from loguru import logger

from quber.agents.llm_client import LLMClient
from quber.core.extractors.base import CellFlag

if TYPE_CHECKING:
    from docling_core.types.doc.document import DoclingDocument

# A real heading is printed once, where its section starts. Two occurrences
# can still be a legitimate continuation heading; from three pages up the
# repetition pattern is layout, not authorship, and the text goes to review.
NOMINATION_PAGES = 3

# How far, as a share of page height, a heading may sit from a demoted running
# header's position band and still be that running header's first half.
SPLIT_HEADER_TOLERANCE = 0.02


async def review_headings(document: "DoclingDocument", llm: LLMClient, source: str) -> List[CellFlag]:
    """Demote confirmed page decoration from the document's heading stream.

    Relabels demoted headings to ``page_header`` in place on ``document`` and
    returns one flag record per demotion. Returns an empty list when nothing
    is nominated, the review backend does not support it, or the agent keeps
    everything.
    """
    heading_pages: Dict[str, Set[int]] = defaultdict(set)
    heading_tops: Dict[str, List[float]] = defaultdict(list)
    heading_count: Dict[str, int] = defaultdict(int)
    page_heights = {no: p.size.height or 1.0 for no, p in (document.pages or {}).items()}
    for item in document.texts:
        if item.label == DocItemLabel.SECTION_HEADER and item.text.strip():
            page = item.prov[0].page_no if item.prov else 0
            heading_pages[item.text.strip()].add(page)
            heading_count[item.text.strip()] += 1
            if item.prov and item.prov[0].bbox is not None:
                # Distance from the page top, 0..1. Docling boxes measure up
                # from the page bottom, so the top edge is height minus t.
                h = page_heights.get(page, 1.0)
                heading_tops[item.text.strip()].append(max(0.0, (h - item.prov[0].bbox.t) / h))

    nominated = {text: pages for text, pages in heading_pages.items() if len(pages) >= NOMINATION_PAGES}
    if not nominated:
        return []

    def position_evidence(text: str) -> str:
        tops = heading_tops.get(text)
        if not tops:
            return "position unknown"
        lo, hi = min(tops), max(tops)
        band = f"{lo:.0%}-{hi:.0%} down the page" if hi - lo > 0.05 else f"{lo:.0%} down the page"
        fixed = "a fixed position" if hi - lo <= 0.05 else "varying positions"
        return f"printed at {fixed}, {band}"

    lines = [
        f'- "{text}" — {heading_count[text]} occurrence(s) as a heading across {len(pages)} pages '
        f"({', '.join(str(p) for p in sorted(pages))}); {position_evidence(text)}"
        for text, pages in sorted(nominated.items(), key=lambda kv: -len(kv[1]))
    ]
    review = await llm.review_headings("Nominated headings of this document:\n\n" + "\n".join(lines))
    if review is None:
        logger.warning(
            "heading review unavailable on this backend; keeping all {} nominated heading(s)",
            len(nominated),
        )
        return []

    demoted = {
        v.text.strip(): v
        for v in review.verdicts
        if v.verdict != "section_heading" and v.text.strip() in nominated
    }
    if not demoted:
        return []

    # Blast radius: how many text items each demoted heading governs, counted
    # before the relabel with the same walk consumers use — every non-heading
    # item inherits the most recent heading above it.
    governed: Dict[str, int] = defaultdict(int)
    current = ""
    for item in document.texts:
        if item.label == DocItemLabel.SECTION_HEADER:
            current = item.text.strip()
        elif current in demoted:
            governed[current] += 1

    flags: List[CellFlag] = []
    for text, verdict in demoted.items():
        relabeled = 0
        for item in document.texts:
            if item.label == DocItemLabel.SECTION_HEADER and item.text.strip() == text:
                # SectionHeaderItem pins its label type, so the relabel goes
                # around assignment validation; the serialized artifact then
                # carries page_header like any other furniture text.
                object.__setattr__(item, "label", DocItemLabel.PAGE_HEADER)
                relabeled += 1
        pages = sorted(nominated[text])
        logger.info(
            "heading demoted ({}): {!r} — a heading on {} pages, governed {} text item(s)",
            verdict.verdict,
            text[:60],
            len(pages),
            governed[text],
        )
        flags.append(
            CellFlag(
                source=source,
                page=pages[0],
                title=text,
                text=text,
                status="heading_demoted",
                note=(
                    f"{verdict.verdict} on {len(pages)} pages "
                    f"({', '.join(str(p) for p in pages)}); relabeled {relabeled} heading "
                    f"item(s) governing {governed[text]} text item(s); {verdict.reason}"
                ),
            )
        )

    # A running header the layout split in two. The parser labels its first
    # half a heading on one page, and that half is also a real heading elsewhere
    # (a company name on the cover and above each statement), so the per-text
    # verdict cannot demote it without demoting the real ones. The occurrence is
    # judged by where it is printed instead: a heading whose text opens a
    # demoted running header, printed where that running header sits, is the
    # running header.
    bands = {
        text: (min(heading_tops[text]), max(heading_tops[text]))
        for text, verdict in demoted.items()
        if verdict.verdict == "running_header" and heading_tops.get(text)
    }
    for item in document.texts:
        if item.label != DocItemLabel.SECTION_HEADER or not item.prov or item.prov[0].bbox is None:
            continue
        text = item.text.strip()
        page = item.prov[0].page_no
        h = page_heights.get(page, 1.0)
        top = max(0.0, (h - item.prov[0].bbox.t) / h)
        for header, (lo, hi) in bands.items():
            if not (header.startswith(text) and len(text) < len(header)):
                continue
            if not (lo - SPLIT_HEADER_TOLERANCE <= top <= hi + SPLIT_HEADER_TOLERANCE):
                continue
            object.__setattr__(item, "label", DocItemLabel.PAGE_HEADER)
            logger.info(
                "heading demoted (split running header): {!r} on page {} at {:.0%} down the page opens {!r}",
                text[:60],
                page,
                top,
                header[:60],
            )
            flags.append(
                CellFlag(
                    source=source,
                    page=page,
                    title=text,
                    text=text,
                    status="heading_demoted",
                    note=(
                        f"first half of the running header {header!r}, printed on page {page} "
                        f"at {top:.0%} down the page where that header sits ({lo:.0%}-{hi:.0%})"
                    ),
                )
            )
            break
    return flags