"""Heading review: the per-occurrence relabel of a running header the layout split."""

from __future__ import annotations

import asyncio
from typing import List, Optional

from docling_core.types.doc.base import BoundingBox, CoordOrigin
from docling_core.types.doc.document import DoclingDocument, PageItem, ProvenanceItem, Size
from docling_core.types.doc.labels import DocItemLabel

from quber.agents.llm_client import HeadingVerdict, LLMHeadingReview, MockLLMClient
from quber.core.fusion.heading_review import review_headings

PAGE_W, PAGE_H = 612.0, 792.0
RUNNING = "KKR Real Estate Finance Trust Inc. Notes to Condensed Consolidated Financial Statements"
COMPANY = "KKR Real Estate Finance Trust Inc."


class Review(MockLLMClient):
    """A review that calls the long banner a running header and keeps the company name."""

    async def review_headings(self, nominated: str) -> Optional[LLMHeadingReview]:
        return LLMHeadingReview(
            verdicts=[
                HeadingVerdict(text=RUNNING, verdict="running_header", reason="fixed position on every page"),
                HeadingVerdict(text=COMPANY, verdict="section_heading", reason="varying positions"),
            ]
        )


def _heading(doc: DoclingDocument, text: str, page: int, top_share: float) -> None:
    top_pt = PAGE_H * (1.0 - top_share)
    bbox = BoundingBox(l=100.0, r=500.0, t=top_pt, b=top_pt - 10.0, coord_origin=CoordOrigin.BOTTOMLEFT)
    doc.add_heading(text=text, prov=ProvenanceItem(page_no=page, bbox=bbox, charspan=(0, len(text))))


def _document() -> DoclingDocument:
    doc = DoclingDocument(name="synthetic")
    for p in range(1, 7):
        doc.pages[p] = PageItem(page_no=p, size=Size(width=PAGE_W, height=PAGE_H))
    # The banner at 3% down on pages 1-3: nominated, demoted.
    for p in (1, 2, 3):
        _heading(doc, RUNNING, p, 0.03)
    # The company name: a real heading on page 4 at 30% down, and the first
    # half of the split banner on page 5 at 3% down. Page 6 holds a real note.
    _heading(doc, COMPANY, 4, 0.30)
    _heading(doc, COMPANY, 5, 0.03)
    _heading(doc, "Note 3. Commercial Real Estate Loans", 6, 0.10)
    return doc


def _labels(doc: DoclingDocument) -> List[tuple[int, str, str]]:
    return [(t.prov[0].page_no, t.text, t.label.value) for t in doc.texts]


def test_split_running_header_demoted_by_position_only_where_it_sits() -> None:
    doc = _document()
    flags = asyncio.run(review_headings(doc, Review(), "doc.pdf"))

    labels = _labels(doc)
    assert (5, COMPANY, DocItemLabel.PAGE_HEADER.value) in labels
    assert (4, COMPANY, DocItemLabel.SECTION_HEADER.value) in labels
    assert (6, "Note 3. Commercial Real Estate Loans", DocItemLabel.SECTION_HEADER.value) in labels
    assert all(label == DocItemLabel.PAGE_HEADER.value for p, text, label in labels if text == RUNNING)
    # One flag for the banner, one for the split occurrence, naming its page.
    split = [f for f in flags if f.text == COMPANY]
    assert (
        len(split) == 1 and split[0].page == 5 and "first half of the running header" in (split[0].note or "")
    )


def test_no_running_header_demoted_means_no_split_relabel() -> None:
    class KeepAll(MockLLMClient):
        async def review_headings(self, nominated: str) -> Optional[LLMHeadingReview]:
            return LLMHeadingReview(
                verdicts=[HeadingVerdict(text=RUNNING, verdict="section_heading", reason="kept")]
            )

    doc = _document()
    flags = asyncio.run(review_headings(doc, KeepAll(), "doc.pdf"))
    assert flags == []
    assert all(label == DocItemLabel.SECTION_HEADER.value for _p, _t, label in _labels(doc))