"""`aia seed` — idempotent: sources, connectors (from code), curated organizations / providers / benchmarks / hardware as entities
with claims attributed to the curated registry source (tier 2, each entry's own `source_url`)."""
from __future__ import annotations

import logging
from typing import Any

from sqlalchemy.ext.asyncio import AsyncConnection

from aiatlas.connectors import registry as connector_registry
from aiatlas.db import execute, fetch_one, jsonb
from aiatlas.ids import new_id
from aiatlas.ontology.benchmarks import metric_bounds, normalize_metric
from aiatlas.ontology.taxonomy import ORG_TYPE_DEFAULT_KIND, normalize_org_kind
from aiatlas.registry import load, org_ref, organizations, provider_ref
from aiatlas.sdk.facts import EntityRef, Facts
from aiatlas.sdk.writer import FactWriter

log = logging.getLogger(__name__)
REGISTRY_SOURCE_KEY = "ai-atlas.registry"


async def seed(conn: AsyncConnection) -> dict[str, Any]:
    out: dict[str, Any] = {}
    out["sources"] = await _seed_sources(conn)
    out["connectors"] = await _seed_connectors(conn)
    reg_source = await fetch_one(conn, "select id, tier from sources where key = :k", k=REGISTRY_SOURCE_KEY)
    assert reg_source
    out["organizations"] = await _seed_organizations(conn, reg_source["id"])
    out["providers"] = await _seed_providers(conn, reg_source["id"])
    out["benchmarks"] = await _seed_benchmarks(conn, reg_source["id"])
    out["hardware"] = await _seed_hardware(conn, reg_source["id"])
    out["domains"] = await _seed_domains(conn)
    return out


async def _seed_sources(conn: AsyncConnection) -> int:
    n = 0
    for s in load("sources"):
        org = organizations().get(s.get("organization", ""))
        org_id = None
        if org:
            row = await fetch_one(conn, "select entity_id from entity_identifiers where scheme = 'registry_org' and value = :v", v=org["key"])
            org_id = row["entity_id"] if row else None
        await execute(conn, """insert into sources (id, key, name, domain, organization_id, tier, kind, category, base_url, rate_limit_per_min, crawl_interval_s, enabled, priority, notes, meta)
                               values (:id, :key, :name, :domain, :org, :tier, :kind, :cat, :base, :rate, :interval, :enabled, :prio, :notes, cast(:meta as jsonb))
                               on conflict (key) do update set name = excluded.name, domain = excluded.domain, organization_id = coalesce(excluded.organization_id, sources.organization_id),
                               tier = excluded.tier, kind = excluded.kind, category = excluded.category, base_url = excluded.base_url, rate_limit_per_min = excluded.rate_limit_per_min,
                               crawl_interval_s = excluded.crawl_interval_s, priority = excluded.priority, notes = excluded.notes, updated_at = now()""",
                      id=new_id("source"), key=s["key"], name=s["name"], domain=s["domain"], org=org_id, tier=s.get("tier", 2), kind=s.get("kind", "website"),
                      cat=s.get("category", "lab"), base=s.get("base_url"), rate=s.get("rate_limit_per_min", 30), interval=s.get("crawl_interval_s", 86400),
                      enabled=s.get("enabled", True), prio=s.get("priority", 2), notes=s.get("notes"), meta=jsonb({k: v for k, v in s.items() if k not in ("key",)}))
        n += 1
    return n


async def _seed_connectors(conn: AsyncConnection) -> int:
    n = 0
    for name, cls in connector_registry().items():
        src = await fetch_one(conn, "select id from sources where key = :k", k=cls.source_key) if cls.source_key else None
        if cls.source_key and not src:
            log.warning("connector source missing in registry/sources.yaml", extra={"connector": name, "source_key": cls.source_key})
        await execute(conn, """insert into connectors (name, source_id, label, description, enabled, priority, interval_seconds, min_interval_seconds, max_interval_seconds,
                               parser_version, rate_limit_per_min, expected_min_records, next_run_at, meta)
                               values (:name, :src, :label, :desc, :enabled, :prio, :interval, :mini, :maxi, :pv, :rate, :emr, now(), cast(:meta as jsonb))
                               on conflict (name) do update set source_id = coalesce(excluded.source_id, connectors.source_id), label = excluded.label, description = excluded.description,
                               priority = excluded.priority, min_interval_seconds = excluded.min_interval_seconds, max_interval_seconds = excluded.max_interval_seconds,
                               parser_version = excluded.parser_version, rate_limit_per_min = excluded.rate_limit_per_min, expected_min_records = excluded.expected_min_records,
                               meta = connectors.meta || excluded.meta, updated_at = now()""",
                      name=name, src=src["id"] if src else None, label=cls.label or name, desc=cls.description or None,
                      enabled=bool(getattr(cls, "enabled_by_default", True)), prio=cls.priority, interval=cls.interval_seconds, mini=cls.min_interval_seconds,
                      maxi=cls.max_interval_seconds, pv=cls.parser_version, rate=cls.rate_per_min, emr=cls.expected_min_records,
                      meta=jsonb({"version": cls.version, "tier": cls.tier, "needs_llm": cls.needs_llm, "module": cls.__module__}))
        n += 1
    return n


def _writer(conn: AsyncConnection, source_id: str, url: str | None) -> FactWriter:
    return FactWriter(conn, source_id=source_id, snapshot_id=None, source_url=url, tier=2, connector_name="registry", extractor="curated", extractor_version="1")


async def _seed_organizations(conn: AsyncConnection, source_id: str) -> int:
    n = 0
    # parents first so `parent` relations resolve
    items = sorted(load("organizations"), key=lambda o: 0 if not o.get("parent") else 1)
    for o in items:
        facts = Facts()
        ref = org_ref(o["key"])
        facts.entities.append(ref)
        for prop in ("country", "headquarters", "founded", "website", "legal_name"):
            if o.get(prop):
                facts.claim(ref, prop, o[prop], source_url=o.get("source_url"))
        if o.get("domains"):
            facts.claim(ref, "domains", o["domains"], source_url=o.get("source_url"))
        if o.get("hf_org"):
            facts.claim(ref, "hf_org", o["hf_org"], source_url=f"https://huggingface.co/{o['hf_org']}")
        if o.get("github_org"):
            facts.claim(ref, "github_org", o["github_org"], source_url=f"https://github.com/{o['github_org']}")
        kind = normalize_org_kind(o.get("kind")) or ORG_TYPE_DEFAULT_KIND.get(o.get("type", "company"))
        if kind:
            facts.claim(ref, "org_kind", kind, source_url=o.get("source_url"))
        if o.get("parent") and o["parent"] in organizations():
            facts.relate(org_ref(o["parent"]), "owns", ref, source_url=o.get("source_url"))
        w = _writer(conn, source_id, o.get("source_url"))
        await w.write(facts)
        n += 1
    return n


async def _seed_providers(conn: AsyncConnection, source_id: str) -> int:
    n = 0
    for p in load("providers"):
        facts = Facts()
        ref = provider_ref(p["key"])
        facts.entities.append(ref)
        for prop in ("website", "pricing_url", "docs_url"):
            if p.get(prop):
                facts.claim(ref, prop, p[prop], source_url=p.get("website"))
        if p.get("organization") in organizations():
            facts.relate(org_ref(p["organization"]), "operates", ref, source_url=p.get("website"))
        await _writer(conn, source_id, p.get("website")).write(facts)
        n += 1
    return n


def benchmark_entity_ref(b: dict[str, Any]) -> EntityRef:
    return EntityRef(entity_type="benchmark", name=b["name"], identifiers={"registry_benchmark": b["key"]}, aliases=list(b.get("aliases", [])), slug_hint=b["key"])


def family_heads(entries: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
    """family key → representative entry (`family_head: true`, else the entry whose key equals the family key)."""
    heads: dict[str, dict[str, Any]] = {}
    for b in entries:
        fam = b.get("family")
        if fam and b.get("family_head"):
            heads[fam] = b
    for b in entries:
        fam = b.get("family")
        if fam and fam not in heads and b["key"] == fam:
            heads[fam] = b
    return heads


async def _seed_benchmarks(conn: AsyncConnection, source_id: str) -> int:
    n = 0
    entries = load("benchmarks")
    heads = family_heads(entries)
    for b in entries:
        facts = Facts()
        ref = benchmark_entity_ref(b)
        facts.entities.append(ref)
        src = b.get("source_url")
        for prop in ("category", "task", "unit", "creator", "website", "paper", "known_limitations", "methodology", "family", "variant", "version", "harness",
                     "comparability_note", "metric_label"):
            if b.get(prop) not in (None, ""):
                facts.claim(ref, prop, b[prop], source_url=src)
        metric = normalize_metric(b.get("metric")) or b.get("metric")
        if metric:
            facts.claim(ref, "metric", metric, source_url=src)
            if b.get("metric") and b["metric"] != metric:
                facts.claim(ref, "metric_raw", b["metric"], source_url=src)
        lo, hi = metric_bounds(metric, b.get("unit"))
        for prop, fallback in (("metric_min", lo), ("metric_max", hi)):
            value = b.get(prop, fallback)
            if value is not None:
                facts.claim(ref, prop, value, source_url=src)
        facts.claim(ref, "higher_is_better", bool(b.get("higher_is_better", True)), source_url=src)
        facts.claim(ref, "family_head", bool(b.get("family_head", False)), source_url=src)
        head = heads.get(b.get("family") or "")
        if head and head["key"] != b["key"]:
            facts.relate(ref, "variant_of", benchmark_entity_ref(head), attributes={"family": b["family"], "variant": b.get("variant")}, source_url=src)
        await _writer(conn, source_id, src).write(facts)
        n += 1
    return n


async def _seed_hardware(conn: AsyncConnection, source_id: str) -> int:
    n = 0
    for h in load("hardware"):
        facts = Facts()
        org = org_ref(h["manufacturer"]) if h.get("manufacturer") in organizations() else None
        ref = EntityRef(entity_type="hardware", name=h["name"], identifiers={"registry_hardware": h["key"]}, aliases=list(h.get("aliases", [])), slug_hint=h["key"], organization=org)
        facts.entities.append(ref)
        for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "tdp_watts", "runtimes", "form_factor", "price_usd"):
            if h.get(prop) not in (None, ""):
                facts.claim(ref, prop, h[prop], source_url=h.get("source_url"))
        facts.claim(ref, "manufacturer", organizations()[h["manufacturer"]]["name"] if org else h.get("manufacturer"), source_url=h.get("source_url"))
        facts.claim(ref, "spec_url", h.get("source_url"), source_url=h.get("source_url"))
        if org:
            facts.relate(org, "manufactures", ref, source_url=h.get("source_url"))
        await _writer(conn, source_id, h.get("source_url")).write(facts)
        n += 1
    return n


async def _seed_domains(conn: AsyncConnection) -> int:
    n = 0
    for o in load("organizations"):
        row = await fetch_one(conn, "select entity_id from entity_identifiers where scheme = 'registry_org' and value = :v", v=o["key"])
        if not row:
            continue
        for d in o.get("domains", []):
            await execute(conn, """insert into domains (domain, organization_id, trust_tier, category) values (:d, :o, 1, :c)
                                   on conflict (domain) do update set organization_id = excluded.organization_id""", d=d, o=row["entity_id"], c=o.get("type", "company"))
            n += 1
    return n


__all__ = ["seed"]