"""/families · /families/{slug} — model families (API 1.1). Canonical `model_family` entities first; pre-canonicalisation data is served
from the legacy `attributes.family` label (items flagged `canonical: false`)."""
from __future__ import annotations

from collections import defaultdict
from typing import Any

from fastapi import APIRouter, Query, Request

from aiatlas.api.common import ENTITY_COLS, ENTITY_FROM, ApiError, cached, entity_summary
from aiatlas.api.detail import LINEAGE_PREDICATES
from aiatlas.db import connection, fetch_all, fetch_one, fetch_val
from aiatlas.ontology.licenses import LICENSES, normalize_license
from aiatlas.services.frontier import all_primary_groups, rank_rows

router = APIRouter(prefix="/api/v1/families", tags=["families"])
MEMBER_SQL = "((f.id is not null and e.family_id = f.id) or (f.id is null and e.attributes->>'family' = f.label))"


def _f(v: Any) -> float | None:
    if v is None or isinstance(v, bool):
        return None
    try:
        return float(v)
    except (TypeError, ValueError):
        return None


def _ranks(groups: dict[str, dict[str, Any]]) -> dict[str, dict[str, dict[str, Any]]]:
    """{model_id: {benchmark slug: {rank, score, metric, config_key}}} over every benchmark's primary group."""
    out: dict[str, dict[str, dict[str, Any]]] = defaultdict(dict)
    for g in groups.values():
        for r in rank_rows(g["rows"], g["higher_is_better"]):
            out[r["model_id"]][g["benchmark"]["slug"]] = {"rank": r["rank"], "score": r["score"], "metric": g["metric"], "config_key": g["config_key"],
                                                          "higher_is_better": g["higher_is_better"], "benchmark": g["benchmark"]["slug"]}
    return out


def _aggregate(members: list[dict[str, Any]], ranks: dict[str, dict[str, dict[str, Any]]]) -> dict[str, Any]:
    params = [p for p in (_f((m["attributes"] or {}).get("parameter_count")) for m in members) if p]
    dates = sorted(str((m["attributes"] or {}).get("release_date") or "")[:10] for m in members if (m["attributes"] or {}).get("release_date"))
    mods: set[str] = set()
    lics: dict[str, int] = defaultdict(int)
    best: dict[str, dict[str, Any]] = {}
    for m in members:
        a = m["attributes"] or {}
        for k in ("modalities", "modalities_input", "modalities_output"):
            if isinstance(a.get(k), list):
                mods |= {str(x).lower() for x in a[k]}
        key = a.get("license_key") or normalize_license(a.get("license"))
        if key or a.get("license"):
            lics[key or str(a.get("license"))] += 1
        for b, rk in ranks.get(m["id"], {}).items():
            if b not in best or rk["rank"] < best[b]["rank"]:
                best[b] = {**rk, "model": m["slug"], "model_name": m["canonical_name"]}
    return {"model_count": len(members), "first_release": dates[0] if dates else None, "last_release": dates[-1] if dates else None,
            "param_range": {"min": min(params), "max": max(params)} if params else None, "modalities": sorted(mods),
            "licenses": [{"key": k, "label": LICENSES[k].label if k in LICENSES else k, "models": n} for k, n in sorted(lics.items(), key=lambda kv: -kv[1])],
            "benchmark_best": dict(sorted(best.items()))}


async def _families(conn: Any) -> list[dict[str, Any]]:
    """Canonical model_family entities + legacy labels not yet backed by an entity."""
    fam = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.entity_type = 'model_family' and e.merged_into is null order by e.canonical_name")
    labels = await fetch_all(conn, """select e.attributes->>'family' as label, min(e.organization_id) as organization_id, count(*) as n from entities e
                                      where e.entity_type = 'model' and e.merged_into is null and e.family_id is null and e.attributes ? 'family' group by 1 order by 1""")
    out = [{"id": r["id"], "slug": r["slug"], "name": r["canonical_name"], "label": r["canonical_name"], "canonical": True, "summary": entity_summary(r), "organization_id": r["organization_id"]} for r in fam]
    names = {x["name"].lower() for x in out}
    for r in labels:
        if r["label"] and r["label"].lower() not in names:
            out.append({"id": None, "slug": r["label"].lower().replace(" ", "-").replace(".", "-"), "name": r["label"], "label": r["label"], "canonical": False, "summary": None, "organization_id": r["organization_id"]})
    return out


@router.get("")
@cached(300)
async def list_families(request: Request, q: str | None = Query(None, max_length=120), org: str | None = None, limit: int = Query(50, ge=1, le=200), offset: int = Query(0, ge=0),
                        sort: str = Query("models", pattern="^(models|name|last_release)$")) -> dict[str, Any]:
    async with connection() as conn:
        fams = await _families(conn)
        members = await fetch_all(conn, f"""select e.id, e.slug, e.canonical_name, e.attributes, e.family_id, e.attributes->>'family' as label, e.organization_id, eo.slug as org_slug, eo.canonical_name as org_name
                                            from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and (e.family_id is not null or e.attributes ? 'family')""")
        groups = await all_primary_groups(conn)
    ranks = _ranks(groups)
    by_fid: dict[str, list[dict[str, Any]]] = defaultdict(list)
    by_label: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for m in members:
        if m["family_id"]:
            by_fid[m["family_id"]].append(m)
        elif m["label"]:
            by_label[m["label"].lower()].append(m)
    items = []
    for f in fams:
        mem = by_fid.get(f["id"], []) if f["id"] else by_label.get(f["name"].lower(), [])
        if not mem and not f["canonical"]:
            continue
        if q and q.lower() not in f["name"].lower():
            continue
        orgs = defaultdict(int)
        for m in mem:
            if m["organization_id"]:
                orgs[(m["organization_id"], m["org_slug"], m["org_name"])] += 1
        top_org = max(orgs.items(), key=lambda kv: kv[1])[0] if orgs else None
        if org and not (top_org and (top_org[1] == org or top_org[0] == org or (top_org[2] or "").lower() == org.lower())):
            continue
        items.append({"id": f["id"], "slug": f["slug"], "name": f["name"], "canonical": f["canonical"], "entity_type": "model_family",
                      "organization": {"id": top_org[0], "slug": top_org[1], "name": top_org[2]} if top_org else None, **_aggregate(mem, ranks)})
    key = {"models": lambda x: (-x["model_count"], x["name"]), "name": lambda x: x["name"].lower(), "last_release": lambda x: (x["last_release"] or "", x["name"])}[sort]
    items.sort(key=key, reverse=(sort == "last_release"))
    return {"items": items[offset:offset + limit], "total": len(items), "limit": limit, "offset": offset,
            "note": "canonical: true = model_family entity; false = legacy attributes.family label awaiting canonicalisation. benchmark_best = best rank of any member in each benchmark's primary group."}


@router.get("/{slug}")
@cached(300)
async def family_detail(request: Request, slug: str, limit: int = Query(100, ge=1, le=500)) -> dict[str, Any]:
    async with connection() as conn:
        fam = await fetch_one(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.entity_type = 'model_family' and e.merged_into is null and (e.slug = :s or e.id = :s or e.canonical_name ilike :s) limit 1", s=slug)
        label = fam["canonical_name"] if fam else slug
        if fam:
            members = await fetch_all(conn, f"select {ENTITY_COLS}, e.family_id from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and (e.family_id = :fid or e.attributes->>'family' ilike :label) "
                                            f"order by e.attributes->>'release_date' asc nulls last, e.canonical_name limit :lim", fid=fam["id"], label=label, lim=limit)
        else:
            members = await fetch_all(conn, f"select {ENTITY_COLS}, e.family_id from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and e.family_id is null "
                                            f"and (e.attributes->>'family' ilike :label or replace(replace(lower(e.attributes->>'family'), ' ', '-'), '.', '-') = lower(:slug)) "
                                            f"order by e.attributes->>'release_date' asc nulls last, e.canonical_name limit :lim", label=label, slug=slug, lim=limit)
        if not fam and not members:
            raise ApiError(404, f"family {slug!r} not found")
        ids = [m["id"] for m in members]
        arts = await fetch_val(conn, "select count(*) from entities where entity_type = 'artifact' and merged_into is null and canonical_id = any(cast(:ids as text[]))", ids=ids) if ids else 0
        providers = await fetch_all(conn, f"select distinct {ENTITY_COLS} from prices p join entities e on e.id = p.provider_id left join entities eo on eo.id = e.organization_id where p.model_id = any(cast(:ids as text[])) and p.valid_to is null order by e.canonical_name", ids=ids) if ids else []
        edges = await fetch_all(conn, "select r.subject_id, r.predicate, r.object_id from relations r where r.valid_to is null and r.predicate = any(cast(:p as text[])) and r.subject_id = any(cast(:ids as text[])) and r.object_id = any(cast(:ids as text[]))",
                                p=[*LINEAGE_PREDICATES, "superseded_by"], ids=ids) if ids else []
        events = await fetch_all(conn, """select ev.entity_id, ev.event_type, ev.summary, ev.occurred_at, ev.importance from change_events ev where ev.entity_id = any(cast(:ids as text[]))
                                          and ev.event_type in ('NEW_MODEL','RELEASE','DEPRECATION_ANNOUNCED','RETIREMENT_ANNOUNCED','STATUS_CHANGED') order by ev.occurred_at asc limit 500""", ids=ids) if ids else []
        groups = await all_primary_groups(conn)
    ranks = _ranks(groups)
    mem_summ = []
    for m in members:
        a = m["attributes"] or {}
        rk = ranks.get(m["id"], {})
        mem_summ.append({"model": entity_summary(m), "key_facts": {k: a.get(k) for k in ("release_date", "parameter_count", "active_parameter_count", "context_length", "openness", "license", "modalities", "status") if a.get(k) not in (None, "", [])},
                         "benchmark_ranks": {b: v["rank"] for b, v in sorted(rk.items())},
                         "benchmark_best": [{"benchmark": b, "rank": v["rank"], "score": v["score"], "metric": v["metric"], "config_key": v["config_key"]} for b, v in sorted(rk.items(), key=lambda kv: kv[1]["rank"])]})
    timeline = [{"date": str((m["attributes"] or {}).get("release_date") or "")[:10] or None, "kind": "release", "model": {"id": m["id"], "slug": m["slug"], "name": m["canonical_name"]}} for m in members]
    timeline += [{"date": e["occurred_at"], "kind": e["event_type"], "summary": e["summary"], "model_id": e["entity_id"]} for e in events if e["event_type"] != "NEW_MODEL"]
    timeline.sort(key=lambda x: str(x["date"] or ""))
    return {"id": fam["id"] if fam else None, "slug": fam["slug"] if fam else slug, "name": label, "canonical": bool(fam), "entity_type": "model_family",
            "summary": entity_summary(fam) if fam else None, **_aggregate(members, ranks), "members": mem_summ, "artifacts_count": int(arts or 0),
            "providers": [entity_summary(p) for p in providers], "lineage": [{"source": e["subject_id"], "target": e["object_id"], "predicate": e["predicate"]} for e in edges],
            "timeline": timeline, "note": None if fam else "served from the legacy attributes.family label (no model_family entity yet)"}