from app.claims import build_claim_sets
from app.models import AtlasAlias, AtlasClaim, AtlasClaimObject, AtlasEntity, AtlasProvenance


def test_build_claim_sets_attaches_provenance_per_claim():
    entity = AtlasEntity(
        atlas_id="atlas:0cqt90abcd123456",
        canonical_label="Donald Trump",
        canonical_description="45th and 47th U.S. President",
        entity_type="Person",
        aliases=[AtlasAlias(label="Donald Trump")],
        claims=[
            AtlasClaim(
                claim_id="clm_raw_ident_mid_/m/0cqt90",
                subject="atlas:0cqt90abcd123456",
                predicate="atlas:hasIdentifier",
                object=AtlasClaimObject(kind="identifier", id_type="mid", value="/m/0cqt90"),
                layer="raw",
                provenance=AtlasProvenance(source="google", retrieval_method="trends-resolution", confidence=0.9, retrieved_at="2026-04-03T00:00:00Z"),
            ),
            AtlasClaim(
                claim_id="clm_raw_ident_qid_Q22686",
                subject="atlas:0cqt90abcd123456",
                predicate="atlas:hasIdentifier",
                object=AtlasClaimObject(kind="identifier", id_type="qid", value="Q22686"),
                layer="raw",
                provenance=AtlasProvenance(source="wikidata", retrieval_method="wbsearchentities + entitydata", confidence=0.99, retrieved_at="2026-04-03T00:00:01Z"),
            ),
            AtlasClaim(
                claim_id="clm_drv_canonical_type",
                subject="atlas:0cqt90abcd123456",
                predicate="atlas:hasCanonicalType",
                object=AtlasClaimObject(kind="type", value="atlas:Person"),
                layer="derived",
                provenance=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"