from __future__ import annotations

from typing import Any

from grantflow.swarm.citations import append_citations
from grantflow.swarm.versioning import append_draft_version

from grantflow.tests.contracts_helpers import fixture_json


def test_append_citations_deduplicates_and_preserves_traceability():
    state = {}
    incoming = [
        {
            "stage": "architect",
            "citation_type": "rag_claim_support",
            "namespace": "usaid_ads201",
            "source": "ads.pdf",
            "page": 12,
            "chunk": 1,
            "chunk_id": "usaid_ads201_p12_c1",
            "used_for": "toc.project_goal",
            "statement_path": "toc.project_goal",
            "label": "USAID ADS 201 p.12",
        },
        {
            "stage": "architect",
            "citation_type": "rag_claim_support",
            "namespace": "usaid_ads201",
            "source": "ads.pdf",
            "page": 12,
            "chunk": 1,
            "chunk_id": "usaid_ads201_p12_c1",
            "used_for": "toc.project_goal",
            "statement_path": "toc.project_goal",
            "label": "USAID ADS 201 p.12",
        },
    ]
    append_citations(state, incoming)
    assert len(state["citations"]) == 1
    assert state["citations"][0]["doc_id"] == "usaid_ads201_p12_c1"
    assert state["citations"][0]["chunk_id"] == "usaid_ads201_p12_c1"


def test_append_citations_prefers_higher_quality_duplicate():
    state = {}
    append_citations(
        state,
        [
            {
                "stage": "architect",
                "citation_type": "rag_low_confidence",
                "namespace": "usaid_ads201",
                "source": "ads.pdf",
                "page": 12,
                "chunk_id": "usaid_ads201_p12_c1",
                "used_for": "toc.project_goal",
                "statement_path": "toc.project_goal",
                "label": "USAID ADS 201 p.12",
                "citation_confidence": 0.2,
                "retrieval_confidence": 0.25,
                "retrieval_rank": 4,
            },
            {
                "stage": "architect",
                "citation_type": "rag_low_confidence",
                "namespace": "usaid_ads201",
                "source": "ads.pdf",
                "page": 12,
                "chunk_id": "usaid_ads201_p12_c1",
                "used_for": "toc.project_goal",
                "statement_path": "toc.project_goal",
                "label": "USAID ADS 201 p.12",
                "citation_confidence": 0.72,
                "retrieval_confidence": 0.82,
                "retrieval_rank": 1,
            },
        ],
    )
    assert len(state["citations"]) == 1
    row = state["citations"][0]
    assert row["citation_confidence"] == 0.72
    assert row["retrieval_confidence"] == 0.82
    assert row["retrieval_rank"] == 1


def test_citations_and_versioning_match_golden_snapshot():
    fixture = fixture_json("citations_versioning_golden.json")
    incoming = fixture.get("incoming_citations") if isinstance(fixture, dict) else None
    expected = fixture.get("expected") if isinstance(fixture, dict) else None
    assert isinstance(incoming, list) and isinstance(expected, dict)

    state: dict[str, Any] = {}
    append_citations(state, incoming)
    append_draft_version(state, section="toc", content={"toc": {"brief": "v1"}}, node="architect", iteration=1)
    append_draft_version(state, section="toc", content={"toc": {"brief": "v1"}}, node="architect", iteration=1)
    append_draft_version(state, section="toc", content={"toc": {"brief": "v2"}}, node="architect", iteration=2)

    actual = {
        "citations": state.get("citations") or [],
        "draft_versions": [
            {
                "version_id": row.get("version_id"),
                "sequence": row.get("sequence"),
                "section": row.get("section"),
                "node": row.get("node"),
                "iteration": row.get("iteration"),
            }
            for row in (state.get("draft_versions") or [])
            if isinstance(row, dict)
        ],
    }
    assert actual == expected


def test_append_draft_version_skips_duplicate_consecutive_payloads():
    state = {}
    toc_payload = {"toc": {"brief": "v1"}}
    append_draft_version(state, section="toc", content=toc_payload, node="architect", iteration=1)
    append_draft_version(state, section="toc", content=toc_payload, node="architect", iteration=1)
    append_draft_version(state, section="toc", content={"toc": {"brief": "v2"}}, node="architect", iteration=2)

    versions = state["draft_versions"]
    assert len(versions) == 2
    assert versions[0]["version_id"] == "toc_v1"
    assert versions[1]["version_id"] == "toc_v2"
    assert versions[0]["sequence"] == 1
    assert versions[1]["sequence"] == 2
