from app.claims import build_claim_sets
from app.models import AtlasAlias, AtlasEntity, AtlasIdentifier, AtlasProvenance


def test_build_claim_sets_attaches_provenance_per_claim():
    entity = AtlasEntity(
        atlas_id="atlas:mid:/m/0cqt90",
        canonical_label="Donald Trump",
        canonical_description="45th and 47th U.S. President",
        entity_type="Person",
        aliases=[AtlasAlias(label="Donald Trump")],
        identifiers=[
            AtlasIdentifier(value="/m/0cqt90", source="google", identifier_type="mid"),
            AtlasIdentifier(value="Q22686", source="wikidata", identifier_type="qid"),
        ],
        provenance=[
            AtlasProvenance(source="google", retrieval_method="trends-resolution", confidence=0.9, retrieved_at="2026-04-03T00:00:00Z"),
            AtlasProvenance(source="wikidata", retrieval_method="wbsearchentities + entitydata", confidence=0.99, retrieved_at="2026-04-03T00:00:01Z"),
            AtlasProvenance(source="openai-llm", retrieval_method="type-classification", confidence=1.0, retrieved_at="2026-04-03T00:00:02Z"),
        ],
        raw_payload={"wikidata": {"status": "ok", "qid": "Q22686", "retrieved_at": "2026-04-03T00:00:01Z"}},
        needs_curation=False,
    )

    raw_claims, derived_claims = build_claim_sets(entity)

    mid_claim = next(c for c in raw_claims if c["object"].get("value") == "/m/0cqt90")
    qid_claim = next(c for c in raw_claims if c["object"].get("value") == "Q22686")
    canonical_type = next(c for c in derived_claims if c["predicate"] == "atlas:hasCanonicalType")

    assert mid_claim["provenance"]["source"] == "google"
    assert qid_claim["provenance"]["source"] == "wikidata"
    assert canonical_type["provenance"]["method"] == "type-classification"