"""Canonicalization steps on synthetic rows, inside a rolled-back transaction on the dev database. Every step must be idempotent:
a second call right after `apply=True` reports zero changes."""
from __future__ import annotations

import uuid
from collections.abc import AsyncIterator
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy.ext.asyncio import AsyncConnection

from aiatlas import db
from aiatlas.db import execute, fetch_all, fetch_one
from aiatlas.sdk.facts import Facts
from aiatlas.sdk.writer import FactWriter
from aiatlas.services import canonical as canon
from aiatlas.services.merge import enforce_current_results, merge_entities


@pytest.fixture
async def conn() -> AsyncIterator[AsyncConnection]:
    async with db.engine().connect() as c:
        trans = await c.begin()
        try:
            yield c
        finally:
            await trans.rollback()
    await db.dispose()


def _tag() -> str:
    return uuid.uuid4().hex[:8]


async def _write(conn: AsyncConnection, facts: Facts, *, tier: int = 2, source_key: str = "", observed_at: datetime | None = None) -> None:
    await FactWriter(conn, source_id=None, snapshot_id=None, source_url="https://example.com/x", tier=tier, connector_name="test", run_id="run_test",
                     source_key=source_key, observed_at=observed_at).write(facts)


async def _run(step, conn: AsyncConnection, apply: bool, *ids: str) -> canon.StepReport:  # type: ignore[no-untyped-def]
    """Run one step scoped to the test's own entities (the dev database is shared with other work)."""
    rep = canon.StepReport(name=step.__name__)
    await step(conn, rep, apply, scope=set(ids) if ids else None)
    return rep


# ---------------------------------------------------------------------------------------------- variants
async def test_variant_folding(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    org = f.entity("company", f"Zorg {t}")
    base = f.entity("model", f"Zeta {t}", organization=org, identifiers={"artificial_analysis": f"zeta-{t}"})
    b = f.entity("benchmark", f"GPQA {t}", identifiers={"registry_benchmark": f"gpqa-{t}"})
    f.result(model=base, benchmark=b, score=70.0, metric="accuracy", unit="%", config={"evaluator": "AA", "index_version": "4.3", "aa_slug": f"zeta-{t}"})
    await _write(conn, f, source_key="artificialanalysis.ai")
    # the variant was created earlier as its own entity (pre-upgrade behaviour): insert it directly
    await execute(conn, """insert into entities (id, entity_type, canonical_name, slug, status) values (:id, 'model', :n, :slug, 'active')""",
                  id=f"model_test{t}v", n=f"zeta-{t}-xhigh", slug=f"zeta-{t}-xhigh")
    await execute(conn, "insert into entity_identifiers (entity_id, scheme, value) values (:e, 'artificial_analysis', :v)", e=f"model_test{t}v", v=f"zeta-{t}-xhigh")
    await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm) values (:e, :a, :n)", e=f"model_test{t}v", a=f"zeta-{t}-xhigh", n=f"zeta{t}xhigh")
    await execute(conn, """insert into benchmark_results (id, model_id, benchmark_id, score, metric, unit, config, dedupe_key, tier)
                           values (:id, :m, :b, 75.0, 'accuracy', '%', cast(:c as jsonb), :d, 2)""",
                  id=f"res_test{t}", m=f"model_test{t}v", b=b.id, c='{"evaluator": "AA", "index_version": "4.3", "aa_slug": "zeta-%s-xhigh"}' % t, d=f"test{t}")
    rep = await _run(canon.step_variants, conn, False, base.id, f"model_test{t}v")
    assert rep.counts.get("folded", 0) >= 1 and not await fetch_one(conn, "select 1 from entities where id = :id and merged_into is not null", id=f"model_test{t}v")
    rep = await _run(canon.step_variants, conn, True, base.id, f"model_test{t}v")
    assert rep.counts["folded"] >= 1
    v = await fetch_one(conn, "select merged_into, status from entities where id = :id", id=f"model_test{t}v")
    assert v["merged_into"] == base.id and v["status"] == "merged"
    rows = await fetch_all(conn, "select score, config, model_id, is_current from benchmark_results where model_id = :m order by score", m=base.id)
    assert [r["score"] for r in rows] == [70.0, 75.0]
    assert rows[1]["config"]["reasoning_effort"] == "xhigh" and rows[1]["config"]["aa_variant_slug"] == f"zeta-{t}-xhigh"
    assert all(r["is_current"] for r in rows)          # same run group, different condition → both current
    assert await fetch_one(conn, "select 1 from resolution_decisions where a_id = :a and b_id = :b and decision = 'merge' and applied", a=f"model_test{t}v", b=base.id)
    again = await _run(canon.step_variants, conn, True, base.id, f"model_test{t}v")
    assert again.counts.get("folded", 0) == 0 and again.counts.get("unresolved_flagged", 0) == 0


async def test_variant_with_official_identifier_is_not_folded(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    org = f.entity("company", f"Zorg {t}")
    f.entity("model", f"Zeta {t}", organization=org, identifiers={"artificial_analysis": f"zeta-{t}"})
    real = f.entity("model", f"Zeta {t} Thinking", organization=org, identifiers={"hf_repo": f"zorg/Zeta-{t}-Thinking"})
    await _write(conn, f)
    rep = await _run(canon.step_variants, conn, True, real.id)
    assert rep.counts.get("_kept_real_model", 0) >= 1
    assert (await fetch_one(conn, "select merged_into from entities where id = :id", id=real.id))["merged_into"] is None


# ---------------------------------------------------------------------------------------------- artifacts
async def test_artifact_detection(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    org = f.entity("company", f"Zorg {t}", attributes={"hf_org": f"zorg{t}"})
    base = f.entity("model", f"Zeta-{t}-9B", organization=org, identifiers={"hf_repo": f"zorg{t}/Zeta-{t}-9B"}, attributes={"hf_repo": f"zorg{t}/Zeta-{t}-9B"})
    third = f.entity("model", f"unsloth/Zeta-{t}-9B-GGUF", organization=f.entity("organization", "unsloth-test", attributes={"hf_org": "unsloth"}),
                     identifiers={"hf_repo": f"unsloth/Zeta-{t}-9B-GGUF"}, attributes={"hf_repo": f"unsloth/Zeta-{t}-9B-GGUF", "quant_format": "gguf", "is_quantized": True})
    official_fp8 = f.entity("model", f"Zeta-{t}-9B-FP8", organization=org, identifiers={"hf_repo": f"zorg{t}/Zeta-{t}-9B-FP8"}, attributes={"hf_repo": f"zorg{t}/Zeta-{t}-9B-FP8"})
    native = f.entity("model", f"Zeta-{t}-Big", organization=org, identifiers={"hf_repo": f"zorg{t}/Zeta-{t}-Big"}, attributes={"hf_repo": f"zorg{t}/Zeta-{t}-Big", "quant_format": "fp8", "is_quantized": True})
    endpoint = f.entity("model", f"Zeta {t} Served", organization=org, identifiers={"openrouter": f"zorg/zeta-{t}-served"}, attributes={"quant_format": "fp8"})
    await _write(conn, f)
    ids = [base.id, third.id, official_fp8.id, native.id, endpoint.id]
    rep = await _run(canon.step_artifacts, conn, True, *ids)
    assert rep.counts.get("artifacts_marked", 0) >= 2
    rows = {r["id"]: r for r in await fetch_all(conn, "select id, entity_type, artifact_kind, canonical_id, identity_confidence from entities where id = any(cast(:ids as text[]))",
                                                    ids=[base.id, third.id, official_fp8.id, native.id, endpoint.id])}
    assert rows[base.id]["entity_type"] == "model"
    assert rows[third.id]["entity_type"] == "artifact" and rows[third.id]["artifact_kind"] == "quantization" and rows[third.id]["canonical_id"] == base.id
    assert rows[official_fp8.id]["entity_type"] == "artifact" and rows[official_fp8.id]["canonical_id"] == base.id       # same org has the plain model
    assert rows[native.id]["entity_type"] == "model"                 # developer repo in native FP8 with no plain sibling = the model
    assert rows[endpoint.id]["entity_type"] == "model"               # provider endpoint precision says nothing about identity
    assert await fetch_one(conn, "select 1 from relations where subject_id = :a and predicate = 'artifact_of' and object_id = :m and valid_to is null", a=third.id, m=base.id)
    # slug, identifiers and claims untouched
    assert (await fetch_one(conn, "select slug from entities where id = :a", a=third.id))["slug"].startswith("unsloth")
    assert await fetch_one(conn, "select 1 from entity_identifiers where entity_id = :a", a=third.id)
    again = await _run(canon.step_artifacts, conn, True, *ids)
    assert again.counts.get("artifacts_marked", 0) == 0 and again.counts.get("artifacts_unresolved", 0) == 0


async def test_unresolved_artifact_gets_review(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    lone = f.entity("model", f"bartowski/Omega-{t}-70B-GGUF", identifiers={"hf_repo": f"bartowski/Omega-{t}-70B-GGUF"}, attributes={"hf_repo": f"bartowski/Omega-{t}-70B-GGUF"})
    await _write(conn, f)
    rep = await _run(canon.step_artifacts, conn, True, lone.id)
    assert rep.counts.get("artifacts_unresolved", 0) >= 1
    row = await fetch_one(conn, "select entity_type, canonical_id, identity_confidence from entities where id = :id", id=lone.id)
    assert row["entity_type"] == "artifact" and row["canonical_id"] is None and row["identity_confidence"] == "low"
    assert await fetch_one(conn, "select 1 from review_queue where kind = 'unresolved_artifact' and :id = any(entity_ids)", id=lone.id)


# ---------------------------------------------------------------------------------------------- families
async def test_families(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    org = f.entity("company", f"Meta test {t}")
    v = str(int(t, 16) % 900 + 100)
    a = f.entity("model", f"Llama {v}.1 9B Test", organization=org)
    b = f.entity("model", f"Llama {v}.1 70B Test", organization=org)
    await _write(conn, f)
    rep = await _run(canon.step_families, conn, True, a.id, b.id)
    assert rep.counts.get("families_created", 0) >= 1 and rep.counts.get("models_linked", 0) >= 2
    rows = await fetch_all(conn, "select e.family_id, fam.canonical_name, fam.entity_type, fam.attributes from entities e join entities fam on fam.id = e.family_id where e.id in (:a, :b)", a=a.id, b=b.id)
    assert len(rows) == 2 and len({r["family_id"] for r in rows}) == 1 and rows[0]["canonical_name"] == f"Llama {v}.1" and rows[0]["entity_type"] == "model_family"
    assert rows[0]["attributes"]["family_root"] == "Llama"
    assert (await fetch_one(conn, "select attributes->>'family' as f from entities where id = :a", a=a.id))["f"] == f"Llama {v}.1"
    again = await _run(canon.step_families, conn, True, a.id, b.id)
    assert again.changes == 0


# ---------------------------------------------------------------------------------------------- licences / openness
async def test_licenses_and_openness(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    m = f.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"hf_repo": f"zorg/zeta-{t}", "license": "Apache 2.0"})
    llama = f.entity("model", f"Zeta {t} L", identifiers={"hf_repo": f"zorg/zeta-{t}-l"}, attributes={"hf_repo": f"zorg/zeta-{t}-l", "license": "llama3.1", "openness": "restricted"})
    closed = f.entity("model", f"Zeta {t} API", attributes={"openness": "proprietary"})
    await _write(conn, f)
    rep = await _run(canon.step_licenses, conn, True, m.id, llama.id, closed.id)
    assert rep.counts.get("uses_license_relations", 0) >= 2
    lic = await fetch_one(conn, "select id, attributes from entities where entity_type = 'license' and slug = 'apache-2.0'")
    assert lic and lic["attributes"]["commercial_use"] is True
    assert await fetch_one(conn, "select 1 from relations where subject_id = :m and predicate = 'uses_license' and object_id = :l", m=m.id, l=lic["id"])
    attrs = (await fetch_one(conn, "select attributes from entities where id = :m", m=m.id))["attributes"]
    assert attrs["openness"] == "open-weights" and attrs["weights_available"] is True and attrs["commercial_use_allowed"] is True and attrs["license_key"] == "Apache-2.0"
    attrs = (await fetch_one(conn, "select attributes from entities where id = :m", m=llama.id))["attributes"]
    assert attrs["openness"] == "restricted-weights" and attrs["openness_raw"] == "restricted" and attrs["license_key"] == "Llama-3.1-Community"
    attrs = (await fetch_one(conn, "select attributes from entities where id = :m", m=closed.id))["attributes"]
    assert attrs["weights_available"] is False and attrs["openness"] == "proprietary"
    claim = await fetch_one(conn, "select tier, extractor, source_id from claims where entity_id = :m and property = 'weights_available' and status = 'current'", m=m.id)
    assert claim["tier"] == 2 and claim["extractor"] == "derived"
    again = await _run(canon.step_licenses, conn, True, m.id, llama.id, closed.id)
    assert again.changes == 0


async def test_derived_openness_never_supersedes_tier1(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    m = f.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"hf_repo": f"zorg/zeta-{t}", "license": "llama3.1", "openness": "open-weights"})
    await _write(conn, f, tier=1)
    await _run(canon.step_licenses, conn, True, m.id)
    cur = await fetch_one(conn, "select value, tier from claims where entity_id = :m and property = 'openness' and status = 'current'", m=m.id)
    assert cur["value"] == "open-weights" and cur["tier"] == 1
    assert await fetch_one(conn, "select 1 from claims where entity_id = :m and property = 'openness' and status = 'conflicting' and extractor = 'derived'", m=m.id)


# ---------------------------------------------------------------------------------------------- taxonomy in place
async def test_taxonomy_in_place(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    hw = f.entity("hardware", f"Box {t}")
    await _write(conn, f)
    # legacy encodings inserted directly (pre-upgrade rows)
    await execute(conn, """insert into claims (id, entity_id, property, value, value_text, tier, status) values (:id, :e, 'kind', '"computer"', 'computer', 1, 'current')""",
                  id=f"claim_test{t}", e=hw.id)
    await execute(conn, """update entities set attributes = attributes || '{"kind": "computer", "status": "limited availability"}' where id = :e""", e=hw.id)
    rep = await _run(canon.step_taxonomy, conn, True, hw.id)
    assert rep.counts.get("claims_normalized:kind", 0) >= 1 and rep.counts.get("attributes_normalized:status", 0) >= 1
    c = await fetch_one(conn, "select value, value_raw from claims where id = :id", id=f"claim_test{t}")
    assert c["value"] == "system" and c["value_raw"] == "computer"
    row = await fetch_one(conn, "select attributes, status from entities where id = :e", e=hw.id)
    assert row["attributes"]["kind"] == "system" and row["attributes"]["kind_raw"] == "computer" and row["attributes"]["status"] == "limited-availability" and row["status"] == "limited-availability"
    assert (await fetch_one(conn, "select canonical from taxonomy_mappings where domain = 'hardware_kind' and raw = 'computer'"))["canonical"] == "system"
    again = await _run(canon.step_taxonomy, conn, True, hw.id)
    assert again.changes == 0


# ---------------------------------------------------------------------------------------------- results
async def test_results_backfill_and_current(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    m = f.entity("model", f"Zeta {t}")
    b = f.entity("benchmark", f"LiveBench {t}", identifiers={"registry_benchmark": f"lb-{t}"})
    await _write(conn, f)
    now = datetime.now(UTC)
    for i, rel in enumerate(["2026-04-01", "2026-05-01", "2026-06-25"]):
        await execute(conn, """insert into benchmark_results (id, model_id, benchmark_id, score, metric, unit, config, dedupe_key, tier, observed_at)
                               values (:id, :m, :b, :s, 'global_average', '%', cast(:c as jsonb), :d, 2, :o)""",
                      id=f"res_t{t}{i}", m=m.id, b=b.id, s=50 + i, c='{"release": "%s"}' % rel, d=f"t{t}{i}", o=now + timedelta(seconds=i))
    rep = await _run(canon.step_results, conn, True, m.id)
    assert rep.counts["results_backfilled"] >= 3 and rep.counts["older_run_rows_closed"] >= 2
    rows = await fetch_all(conn, "select score, is_current, valid_to, run_group, config_key, trust_level from benchmark_results where model_id = :m order by score", m=m.id)
    assert [(r["score"], r["is_current"], r["valid_to"] is None) for r in rows] == [(50, False, False), (51, False, False), (52, True, True)]
    assert len({r["config_key"] for r in rows}) == 1 and rows[0]["run_group"] == "2026-04-01" and rows[0]["trust_level"] == "unverified"
    assert await enforce_current_results(conn, model_id=m.id, dry_run=True) == 0
    again = await _run(canon.step_results, conn, True, m.id)
    assert again.changes == 0


# ---------------------------------------------------------------------------------------------- events
async def test_events_backfill_classification(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    m = f.entity("model", f"Zeta {t}")
    await _write(conn, f, observed_at=datetime.now(UTC))
    old = datetime.now(UTC) - timedelta(days=30)
    await execute(conn, """insert into change_events (id, entity_id, event_type, category, summary, importance, observed_at, effective_at, connector_name, dedupe_key)
                           values (:id, :e, 'RELEASE', 'model', 'old release', 2, now(), :eff, 'curation', :d)""", id=f"evt_t{t}", e=m.id, eff=old, d=f"t{t}")
    rep = await _run(canon.step_events, conn, True, m.id)
    assert rep.counts.get("backfill_flag_set", 0) >= 1
    ev = await fetch_one(conn, "select is_backfill, group_key from change_events where id = :id", id=f"evt_t{t}")
    assert ev["is_backfill"] is True and ev["group_key"] == f"release:{m.id}:{old:%Y-%m}"
    again = await _run(canon.step_events, conn, True, m.id)
    assert again.changes == 0


# ---------------------------------------------------------------------------------------------- merge modes & duplicates
async def test_merge_modes(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    fam = f.entity("model_family", f"Zeta fam {t}")
    a = f.entity("model", f"Zeta {t}")
    b = f.entity("model", f"Zeta {t} dup")
    c = f.entity("model", f"Zeta {t} FP8")
    await _write(conn, f)
    res = await merge_entities(conn, c.id, a.id, mode="variant", payload={"artifact_kind": "conversion"})
    row = await fetch_one(conn, "select entity_type, canonical_id, merged_into from entities where id = :id", id=c.id)
    assert res["mode"] == "variant" and row["entity_type"] == "artifact" and row["canonical_id"] == a.id and row["merged_into"] is None
    await merge_entities(conn, a.id, fam.id, mode="family_member")
    assert (await fetch_one(conn, "select family_id from entities where id = :id", id=a.id))["family_id"] == fam.id
    res = await merge_entities(conn, b.id, a.id, mode="alias")
    row = await fetch_one(conn, "select merged_into, status from entities where id = :id", id=b.id)
    assert row["merged_into"] == a.id and row["status"] == "merged"
    assert await fetch_one(conn, "select 1 from resolution_decisions where a_id = :a and b_id = :b and decision = 'alias'", a=b.id, b=a.id)
    assert await fetch_one(conn, "select 1 from admin_audit_log where action = 'entity.alias' and target = :t", t=a.id)
    with pytest.raises(ValueError):
        await merge_entities(conn, b.id, a.id)          # already merged


async def test_duplicates_step(conn: AsyncConnection) -> None:
    t = _tag()
    f = Facts()
    keep = f.entity("company", f"Kwai {t}", attributes={"hf_org": f"kwai{t}"})
    dup = f.entity("organization", f"kwai{t}", attributes={"hf_org": f"kwai{t}"})
    x = f.entity("model", f"Dup {t}", identifiers={"openai_model_id": f"dup-{t}-a"})
    y = f.entity("model", f"Dup {t} ", identifiers={"openai_model_id": f"dup-{t}-b"})      # same name, conflicting identifier scheme → review only
    await _write(conn, f)
    rep = await _run(canon.step_duplicates, conn, True, keep.id, dup.id, x.id, y.id)
    assert rep.counts.get("merged", 0) >= 1 and rep.counts.get("review_conflicting_identifiers", 0) >= 1
    assert (await fetch_one(conn, "select merged_into from entities where id = :id", id=dup.id))["merged_into"] == keep.id
    assert (await fetch_one(conn, "select attributes->>'org_kind' as k from entities where id = :id", id=keep.id))["k"] == "company"
    for i in (x.id, y.id):
        assert (await fetch_one(conn, "select merged_into from entities where id = :id", id=i))["merged_into"] is None
    again = await _run(canon.step_duplicates, conn, True, keep.id, dup.id, x.id, y.id)
    assert again.counts.get("merged", 0) == 0