"""Benchmark comparability groups, per-model leaderboards and the "frontier" model set (API 1.1). Deterministic, no LLM.

A *group* is (benchmark, canonical metric, config_key): the comparability-relevant part of the result configuration hashed by
`ontology.benchmarks.config_key` (task keys + metric). Rows written before migration 0003 carry a NULL `config_key` / `trust_level`;
both are recomputed here from `config` and the source key so the API behaves the same before and after canonicalisation.

Leaderboards are ONE row per canonical model: the best current row of that model inside the chosen group. Effort variants folded into
a model land in the same group (reasoning effort is a *condition*, not a task key) and are summarised in `config`."""
from __future__ import annotations

from collections import Counter, defaultdict
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy.ext.asyncio import AsyncConnection

from aiatlas.db import fetch_all, fetch_val
from aiatlas.ontology.benchmarks import (
    CONDITION_KEYS,
    METRICS,
    TASK_KEYS,
    TRUST_LABELS,
    comparability,
    config_key,
    family_of,
    normalize_metric,
    trust_level,
)
from aiatlas.services import cache

RESULT_SELECT = """
    select r.id, r.model_id, r.benchmark_id, r.score, r.metric, r.unit, r.higher_is_better, r.config, r.evaluated_at, r.observed_at, r.source_url, r.tier,
           r.confidence, r.valid_to, r.config_key, r.trust_level, r.variant, r.run_group, r.is_current, r.extractor, s.key as source_key,
           m.slug as model_slug, m.canonical_name as model_name, m.entity_type as model_type, m.merged_into as model_merged_into, m.organization_id,
           m.attributes as model_attrs, m.family_id, mo.slug as org_slug, mo.canonical_name as org_name,
           b.slug as benchmark_slug, b.canonical_name as benchmark_name, b.attributes as benchmark_attrs
    from benchmark_results r
    join entities m on m.id = r.model_id left join entities mo on mo.id = m.organization_id
    join entities b on b.id = r.benchmark_id
    left join sources s on s.id = r.source_id
"""
CURRENT = "r.valid_to is null and r.is_current"
CANONICAL_MODEL = "m.entity_type = 'model' and m.merged_into is null"
FRONTIER_CACHE_KEY = "frontier:model-ids:v1"
FRONTIER_TTL_S = 600
FRONTIER_METHODOLOGY = ("Frontier models = canonical models released in the last 12 months by organizations with at least 3 canonical models, "
                        "OR holding a top-10 rank in the primary comparability group of at least one benchmark. Artifacts and folded variants are excluded. "
                        "No composite score is used to pick them.")


# ------------------------------------------------------------------------------------------------------------------ row enrichment


def enrich(row: dict[str, Any]) -> dict[str, Any]:
    """Fill `config_key` / `trust_level` / `metric_canonical` when the writer has not (pre-0003 rows)."""
    cfg = row.get("config") or {}
    metric = normalize_metric(row.get("metric")) or (row.get("metric") or "").strip().lower() or "score"
    row["metric_canonical"] = metric
    if not row.get("config_key"):
        row["config_key"] = config_key(cfg, row.get("metric"))
    if not row.get("trust_level"):
        row["trust_level"] = trust_level(row.get("source_key"), cfg, extractor=row.get("extractor") or "deterministic")
    return row


def group_label(metric: str, cfg: dict[str, Any] | None) -> str:
    parts = [f"{k}={cfg[k]}" for k in TASK_KEYS if cfg and cfg.get(k) not in (None, "", [], {})]
    return metric + (" · " + " · ".join(str(p) for p in parts) if parts else "")


def config_summary(cfg: dict[str, Any] | None) -> dict[str, Any]:
    """Task + condition keys only (bookkeeping keys such as aa_slug are dropped)."""
    return {k: cfg[k] for k in (*TASK_KEYS, *CONDITION_KEYS) if cfg and cfg.get(k) not in (None, "", [], {})}


def direction_of(rows: list[dict[str, Any]], metric: str) -> bool:
    votes = Counter(bool(r.get("higher_is_better")) for r in rows if r.get("higher_is_better") is not None)
    if votes:
        return votes.most_common(1)[0][0]
    spec = METRICS.get(metric)
    return bool(spec["higher_is_better"]) if spec else True


# ------------------------------------------------------------------------------------------------------------------ loading


async def load_results(conn: AsyncConnection, *, benchmark_ids: list[str] | None = None, model_ids: list[str] | None = None, current_only: bool = True,
                       canonical_models_only: bool = True, limit: int = 50_000) -> list[dict[str, Any]]:
    where = []
    params: dict[str, Any] = {"lim": limit}
    if benchmark_ids is not None:
        where.append("r.benchmark_id = any(cast(:bids as text[]))")
        params["bids"] = benchmark_ids
    if model_ids is not None:
        where.append("r.model_id = any(cast(:mids as text[]))")
        params["mids"] = model_ids
    if current_only:
        where.append(CURRENT)
    if canonical_models_only:
        where.append(CANONICAL_MODEL)
    sql = RESULT_SELECT + " where " + (" and ".join(where) or "true") + " order by r.benchmark_id, r.observed_at desc limit :lim"
    return [enrich(r) for r in await fetch_all(conn, sql, **params)]


# ------------------------------------------------------------------------------------------------------------------ grouping


def group_rows(rows: list[dict[str, Any]]) -> dict[tuple[str, str, str], dict[str, Any]]:
    """{(benchmark_id, metric, config_key): {rows, metric, config_key, label, n, models, higher_is_better, representative config}}."""
    groups: dict[tuple[str, str, str], dict[str, Any]] = {}
    for r in rows:
        k = (r["benchmark_id"], r["metric_canonical"], r["config_key"])
        g = groups.get(k)
        if g is None:
            g = groups[k] = {"benchmark_id": r["benchmark_id"], "metric": r["metric_canonical"], "config_key": r["config_key"], "rows": [], "models": set(),
                             "config": {k2: r["config"][k2] for k2 in TASK_KEYS if (r.get("config") or {}).get(k2) not in (None, "", [], {})}}
        g["rows"].append(r)
        g["models"].add(r["model_id"])
    for g in groups.values():
        g["n"] = len(g["rows"])
        g["model_count"] = len(g["models"])
        g["label"] = group_label(g["metric"], g["config"])
        g["higher_is_better"] = direction_of(g["rows"], g["metric"])
        g["trust_mix"] = dict(Counter(r["trust_level"] for r in g["rows"]))
    return groups


def primary_metric(benchmark_attrs: dict[str, Any] | None, groups: list[dict[str, Any]]) -> str | None:
    """Registry metric when it has current rows, else the most populated metric that is not a LiveBench per-category average."""
    declared = normalize_metric((benchmark_attrs or {}).get("metric"))
    present = Counter()
    for g in groups:
        present[g["metric"]] += g["n"]
    if declared and present.get(declared):
        return declared
    ranked = [m for m, _ in present.most_common() if not m.startswith("category:")] or [m for m, _ in present.most_common()]
    return ranked[0] if ranked else declared


def primary_group(benchmark_attrs: dict[str, Any] | None, groups: list[dict[str, Any]]) -> dict[str, Any] | None:
    metric = primary_metric(benchmark_attrs, groups)
    cands = [g for g in groups if g["metric"] == metric]
    if not cands:
        return None
    return max(cands, key=lambda g: (g["model_count"], g["n"], g["config_key"]))


def best_per_model(rows: list[dict[str, Any]], higher_is_better: bool) -> list[dict[str, Any]]:
    best: dict[str, dict[str, Any]] = {}
    for r in rows:
        cur = best.get(r["model_id"])
        if cur is None or (r["score"] > cur["score"] if higher_is_better else r["score"] < cur["score"]) or \
           (r["score"] == cur["score"] and (r.get("evaluated_at") or r["observed_at"]) > (cur.get("evaluated_at") or cur["observed_at"])):
            best[r["model_id"]] = r
    return sorted(best.values(), key=lambda r: (-r["score"] if higher_is_better else r["score"], r["model_name"] or ""))


def rank_rows(rows: list[dict[str, Any]], higher_is_better: bool) -> list[dict[str, Any]]:
    """Competition ranking (1, 2, 2, 4) over best-per-model rows."""
    ranked = best_per_model(rows, higher_is_better)
    out: list[dict[str, Any]] = []
    prev_score, prev_rank = None, 0
    for i, r in enumerate(ranked, start=1):
        rank = prev_rank if prev_score is not None and r["score"] == prev_score else i
        prev_score, prev_rank = r["score"], rank
        out.append({**r, "rank": rank})
    return out


def leaderboard_rows(group: dict[str, Any], *, history_rows: list[dict[str, Any]] | None = None, comparable_only: bool = False) -> list[dict[str, Any]]:
    """Public leaderboard rows for one group. `history_rows` (closed rows of the same group) give the previous rank per model."""
    hib = group["higher_is_better"]
    ranked = rank_rows(group["rows"], hib)
    leader = ranked[0] if ranked else None
    prev_rank: dict[str, int] = {}
    if history_rows:
        prev_rank = {r["model_id"]: r["rank"] for r in rank_rows(history_rows, hib)}
    out: list[dict[str, Any]] = []
    for r in ranked:
        level, reasons = comparability(leader["config"] if leader else None, r.get("config"), leader["metric"] if leader else None, r.get("metric")) if leader else ("comparable", [])
        if comparable_only and level != "comparable":
            continue
        prev = prev_rank.get(r["model_id"])
        out.append({
            "rank": r["rank"], "model": _model_ref(r), "score": r["score"], "metric": r["metric_canonical"], "unit": r.get("unit"), "higher_is_better": hib,
            "delta_rank": (prev - r["rank"]) if prev is not None else None, "previous_rank": prev,
            "trust_level": r["trust_level"], "trust_label": TRUST_LABELS.get(r["trust_level"], r["trust_level"]), "config": config_summary(r.get("config")),
            "config_key": r["config_key"], "comparability": level, "comparability_reasons": reasons, "evaluated_at": r.get("evaluated_at"), "observed_at": r["observed_at"],
            "source_url": r.get("source_url"), "tier": r.get("tier"), "result_id": r["id"], "n_rows": sum(1 for x in group["rows"] if x["model_id"] == r["model_id"]),
        })
    return out


def _model_ref(r: dict[str, Any]) -> dict[str, Any]:
    attrs = r.get("model_attrs") or {}
    return {"id": r["model_id"], "slug": r["model_slug"], "name": r["model_name"], "entity_type": r.get("model_type", "model"),
            "organization": {"id": r.get("organization_id"), "slug": r.get("org_slug"), "name": r.get("org_name")} if r.get("organization_id") else None,
            "attributes": {k: attrs[k] for k in ("openness", "parameter_count", "context_length", "release_date", "modalities", "license", "family") if attrs.get(k) not in (None, "", [])}}


def group_summary(g: dict[str, Any]) -> dict[str, Any]:
    return {"metric": g["metric"], "config_key": g["config_key"], "label": g["label"], "n": g["n"], "model_count": g["model_count"], "config": g["config"],
            "higher_is_better": g["higher_is_better"], "trust_mix": g["trust_mix"]}


def frontier_series(rows: list[dict[str, Any]], higher_is_better: bool) -> list[dict[str, Any]]:
    """History of the leader: each time a new best score appears (ordered by coalesce(evaluated_at, observed_at))."""
    ordered = sorted(rows, key=lambda r: ((r.get("evaluated_at") or r["observed_at"]), r["observed_at"], r["id"]))
    best: float | None = None
    out: list[dict[str, Any]] = []
    for r in ordered:
        s = float(r["score"])
        better = best is None or (s > best if higher_is_better else s < best)
        if better:
            best = s
            out.append({"date": (r.get("evaluated_at") or r["observed_at"]), "model": _model_ref(r), "score": s, "trust_level": r["trust_level"],
                        "config": config_summary(r.get("config")), "result_id": r["id"]})
    return out


def leader_at(rows: list[dict[str, Any]], at: datetime, benchmark_attrs: dict[str, Any] | None) -> dict[str, Any] | None:
    """Leader of the primary group as it was known at `at`: rows observed at or before `at` and not closed before `at`."""
    visible = [r for r in rows if r["observed_at"] <= at and (r.get("valid_to") is None or r["valid_to"] > at)]
    if not visible:
        return None
    groups = list(group_rows(visible).values())
    pg = primary_group(benchmark_attrs, groups)
    if not pg:
        return None
    ranked = rank_rows(pg["rows"], pg["higher_is_better"])
    if not ranked:
        return None
    top = ranked[0]
    return {"model": _model_ref(top), "score": top["score"], "metric": pg["metric"], "config_key": pg["config_key"], "group_label": pg["label"],
            "trust_level": top["trust_level"], "n_models": pg["model_count"], "as_of": at}


# ------------------------------------------------------------------------------------------------------------------ benchmark catalogue


async def benchmark_meta(conn: AsyncConnection, ids: list[str] | None = None) -> dict[str, dict[str, Any]]:
    where = "e.entity_type = 'benchmark' and e.merged_into is null" + (" and e.id = any(cast(:ids as text[]))" if ids is not None else "")
    rows = await fetch_all(conn, f"select e.id, e.slug, e.canonical_name, e.attributes from entities e where {where}", ids=ids)
    out: dict[str, dict[str, Any]] = {}
    for r in rows:
        attrs = r["attributes"] or {}
        fam, variant = family_of(r["slug"])
        out[r["id"]] = {"id": r["id"], "slug": r["slug"], "name": r["canonical_name"], "attributes": attrs, "category": attrs.get("category"),
                        "family": attrs.get("family") or fam, "variant": attrs.get("variant") or variant,
                        "metric": normalize_metric(attrs.get("metric")) or attrs.get("metric"), "unit": attrs.get("unit"),
                        "direction": attrs.get("direction") or ("higher" if METRICS.get(normalize_metric(attrs.get("metric")) or "", {}).get("higher_is_better", True) else "lower")}
    return out


async def all_primary_groups(conn: AsyncConnection, *, min_results: int = 1) -> dict[str, dict[str, Any]]:
    """{benchmark_id: primary group (with rows)} across every benchmark, from current rows of canonical models."""
    meta = await benchmark_meta(conn)
    rows = await load_results(conn)
    by_bench: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for r in rows:
        by_bench[r["benchmark_id"]].append(r)
    out: dict[str, dict[str, Any]] = {}
    for bid, brows in by_bench.items():
        groups = list(group_rows(brows).values())
        pg = primary_group(meta.get(bid, {}).get("attributes"), groups)
        if pg and pg["n"] >= min_results:
            pg = dict(pg)
            pg["benchmark"] = meta.get(bid) or {"id": bid, "slug": brows[0]["benchmark_slug"], "name": brows[0]["benchmark_name"]}
            pg["all_groups"] = groups
            out[bid] = pg
    return out


# ------------------------------------------------------------------------------------------------------------------ frontier model set


async def frontier_model_ids(conn: AsyncConnection, *, use_cache: bool = True) -> tuple[set[str], dict[str, Any]]:
    """Frontier composition (see FRONTIER_METHODOLOGY). Returns (ids, sample counts)."""
    if use_cache:
        hit = await cache.cache_get(FRONTIER_CACHE_KEY)
        if hit and hit.get("ids"):
            # guard: the cached set must still be live canonical models of THIS database (cache keys are namespaced per DB, but a
            # canonicalisation run can merge ids away between refreshes) — otherwise recompute
            live = await fetch_val(conn, "select count(*) from entities where id = any(cast(:ids as text[])) and entity_type = 'model' and merged_into is null", ids=hit["ids"])
            if int(live or 0) == len(hit["ids"]):
                return set(hit["ids"]), hit["composition"]
    since = (datetime.now(UTC) - timedelta(days=365)).date().isoformat()
    recent = await fetch_all(conn, """
        select e.id from entities e where e.entity_type = 'model' and e.merged_into is null and e.attributes->>'release_date' >= :since
        and e.organization_id in (select organization_id from entities where entity_type = 'model' and merged_into is null and organization_id is not null
                                  group by organization_id having count(*) >= 3)""", since=since)
    recent_ids = {r["id"] for r in recent}
    top10: set[str] = set()
    groups = await all_primary_groups(conn)
    for g in groups.values():
        for r in rank_rows(g["rows"], g["higher_is_better"]):
            if r["rank"] <= 10:
                top10.add(r["model_id"])
    ids = recent_ids | top10
    composition = {"recent_by_active_orgs": len(recent_ids), "top10_on_a_benchmark": len(top10), "total": len(ids), "since": since}
    await cache.cache_set(FRONTIER_CACHE_KEY, {"ids": sorted(ids), "composition": composition}, FRONTIER_TTL_S)
    return ids, composition


__all__ = ["CANONICAL_MODEL", "CURRENT", "FRONTIER_METHODOLOGY", "RESULT_SELECT", "all_primary_groups", "benchmark_meta", "best_per_model", "config_summary",
           "direction_of", "enrich", "frontier_model_ids", "frontier_series", "group_label", "group_rows", "group_summary", "leader_at", "leaderboard_rows",
           "load_results", "primary_group", "primary_metric", "rank_rows"]