from haiku.rag.capabilities.ledger import (
    CapabilityEvidenceRecord,
    CitationDeclaration,
    EvidenceRef,
    citation_status,
)


def rag_ref(chunk_id: str = "c1") -> EvidenceRef:
    return EvidenceRef(capability="rag", chunk_id=chunk_id)


def test_a_record_survives_the_state_round_trip():
    """Capability state is persisted as JSON, so the schema must survive it.

    A dict keyed by ``(capability, chunk_id)`` does not: the key serialises to
    ``"rag,c1"`` and fails revalidation as a tuple.
    """
    record = CapabilityEvidenceRecord(question=4)
    record.note_evidence(5)
    record.declare([rag_ref()], epoch=7, retrieved_now={"c1"})

    restored = CapabilityEvidenceRecord.model_validate(record.model_dump(mode="json"))

    assert restored == record
    assert citation_status([restored], question=4) == "grounded"
    assert restored.occurrences["c1"].cited_in_questions == [4]
    assert restored.occurrences["c1"].retrieved_in_questions == [4]


def test_no_declaration_reads_as_missing():
    assert citation_status([CapabilityEvidenceRecord()], question=0) == "missing"
    assert citation_status([], question=0) == "missing"


def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded():
    grounded = CapabilityEvidenceRecord()
    grounded.declare([rag_ref()], epoch=1)

    ungrounded = CapabilityEvidenceRecord()
    ungrounded.declare([], epoch=1)

    assert citation_status([grounded], question=0) == "grounded"
    assert citation_status([ungrounded], question=0) == "ungrounded"


def test_an_earlier_questions_declaration_is_never_current():
    """Epochs outlive a question, so the epoch alone would inherit it."""
    record = CapabilityEvidenceRecord(question=2)
    record.declare([rag_ref()], epoch=3)
    assert citation_status([record], question=2) == "grounded"

    record.question = 8

    assert record.declaration is not None
    assert citation_status([record], question=8) == "missing"


def test_a_citation_in_the_same_request_as_the_evidence_is_not_current():
    """Citing must follow seeing: equal epochs mean one request."""
    record = CapabilityEvidenceRecord()
    record.note_evidence(5)
    record.declare([rag_ref()], epoch=5)

    assert citation_status([record], question=0) == "missing"

    record.declare([rag_ref()], epoch=7)

    assert citation_status([record], question=0) == "grounded"


def test_evidence_from_another_capability_after_citing_makes_it_uncited():
    """Currency spans capabilities, which only works because epochs are global."""
    cited = CapabilityEvidenceRecord()
    cited.note_evidence(3)
    cited.declare([rag_ref()], epoch=5)
    searched_after = CapabilityEvidenceRecord()
    searched_after.note_evidence(7)

    assert citation_status([cited], question=0) == "grounded"
    assert citation_status([cited, searched_after], question=0) == "missing"


def test_declarations_at_the_same_epoch_merge():
    record = CapabilityEvidenceRecord()
    record.declare([rag_ref("c1")], epoch=3)
    record.declare([rag_ref("c2")], epoch=3)

    assert record.declaration is not None
    assert [ref.chunk_id for ref in record.declaration.refs] == ["c1", "c2"]


def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it():
    record = CapabilityEvidenceRecord()
    record.declare([rag_ref()], epoch=3)
    record.declare([rag_ref()], epoch=3)

    assert record.declaration is not None
    assert len(record.declaration.refs) == 1


def test_neither_cite_order_downgrades_a_grounded_declaration():
    grounded_then_empty = CapabilityEvidenceRecord()
    grounded_then_empty.declare([rag_ref()], epoch=3)
    grounded_then_empty.declare([], epoch=3)

    empty_then_grounded = CapabilityEvidenceRecord()
    empty_then_grounded.declare([], epoch=3)
    empty_then_grounded.declare([rag_ref()], epoch=3)

    assert citation_status([grounded_then_empty], question=0) == "grounded"
    assert citation_status([empty_then_grounded], question=0) == "grounded"


def test_the_same_chunk_id_under_two_capabilities_stays_separate():
    rag = CapabilityEvidenceRecord()
    rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3)
    analysis = CapabilityEvidenceRecord()
    analysis.declare([EvidenceRef(capability="analysis", chunk_id="shared")], epoch=3)

    assert rag.occurrences["shared"].capability == "rag"
    assert analysis.occurrences["shared"].capability == "analysis"


def test_an_evidence_epoch_never_moves_backwards():
    record = CapabilityEvidenceRecord()
    record.note_evidence(9)
    record.note_evidence(4)

    assert record.latest_evidence_epoch == 9


def test_citing_the_same_chunk_in_two_questions_records_both():
    record = CapabilityEvidenceRecord(question=2)
    record.declare([rag_ref()], epoch=3, retrieved_now={"c1"})
    record.question = 8
    record.declare([rag_ref()], epoch=9)

    occurrence = record.occurrences["c1"]
    assert occurrence.cited_in_questions == [2, 8]
    assert occurrence.retrieved_in_questions == [2]


def test_a_declaration_records_the_question_and_epoch_it_was_made_at():
    record = CapabilityEvidenceRecord(question=6)
    record.declare([rag_ref()], epoch=11)

    assert record.declaration == CitationDeclaration(
        question=6, epoch=11, refs=[rag_ref()]
    )