"""The document identity the routes state ahead of the context.

The block is built from the documents row and prepended to the user message
on every answer path; the system prompt is untouched. Fields the row lacks
are left out rather than printed empty."""

from quber.playground.agent import _prompt as grounded_prompt
from quber.playground.answers import declared, fixed
from quber.playground.answers.document import DocumentIdentity, with_document
from quber.playground.retrieval import RetrievedChunk

ROW = {
    "id": 291,
    "doc_key": "705c9a88ad5d2694",
    "title": "Q1 2026 Earnings Supplement",
    "folder": "KREF",
    "filing_type": "99-2",
    "year": 2026,
    "period": "Q1",
}

CHUNK = RetrievedChunk(
    chunk_id="#/tables/13", chunk_type="table", page=15, content="<table></table>", score=0.01
)


def test_header_lists_the_row_fields_in_order():
    header = DocumentIdentity.from_row(ROW).header()
    assert header.splitlines()[:5] == [
        "Document",
        "  Title:  Q1 2026 Earnings Supplement",
        "  Issuer:  KREF",
        "  Form:  99-2",
        "  Period:  Q1 2026",
    ]
    assert header.endswith("unless a chunk states another period.")


def test_missing_fields_are_left_out():
    header = DocumentIdentity.from_row({"title": "3M 2018 Form 10-K"}).header()
    assert "Issuer" not in header
    assert "Period" not in header
    assert "  Title:  3M 2018 Form 10-K" in header


def test_no_document_leaves_the_prompt_unchanged():
    assert with_document("Context chunks:", None) == "Context chunks:"


def test_every_answer_path_opens_with_the_document():
    identity = DocumentIdentity.from_row(ROW)
    for build in (declared.prompt_for, fixed.prompt_for, grounded_prompt):
        prompt = build("What was book value per share?", [CHUNK], identity)
        assert prompt.startswith("Document\n  Title:  Q1 2026 Earnings Supplement")
        assert "Context chunks:" in prompt
        assert prompt.index("Document") < prompt.index("Context chunks:")


def test_paths_still_work_without_a_document():
    for build in (declared.prompt_for, fixed.prompt_for, grounded_prompt):
        prompt = build("What was book value per share?", [CHUNK])
        assert prompt.startswith("Context chunks:")