"""/compare?ids=a,b,… — side-by-side of 2–6 entities of the same type with per-type dimension sets and provenance."""
from __future__ import annotations

from typing import Any

from fastapi import APIRouter, Query, Request

from aiatlas.api.common import (
    COMPANY_TYPES,
    PRICE_COLS,
    PRICE_FROM,
    RESULT_COLS,
    RESULT_FROM,
    ApiError,
    cached,
    csv,
    enrich_provenance,
    entity_summary,
    price_row,
    resolve_entity,
    result_row,
)
from aiatlas.db import connection, fetch_all

router = APIRouter(prefix="/api/v1/compare", tags=["compare"])

D = lambda key, label, kind="text", unit=None, source="attr": {"key": key, "label": label, "kind": kind, **({"unit": unit} if unit else {}), "source": source}

DIMENSIONS: dict[str, list[dict[str, Any]]] = {
    "model": [D("parameter_count", "Parameters", "number", "params"), D("active_parameter_count", "Active parameters", "number", "params"),
              D("context_length", "Context window", "number", "tokens"), D("max_output_tokens", "Max output", "number", "tokens"), D("openness", "Openness"),
              D("license", "License"), D("modalities", "Modalities", "list"), D("release_date", "Release date", "date"), D("knowledge_cutoff", "Knowledge cutoff", "date"),
              D("status", "Status"), D("family", "Family"), D("architecture", "Architecture"),
              D("best_input_per_mtok", "Best input price", "number", "USD / 1M tokens", "prices"), D("best_output_per_mtok", "Best output price", "number", "USD / 1M tokens", "prices"),
              D("provider_count", "Providers", "number", None, "prices")],
    "provider": [D("website", "Website"), D("pricing_url", "Pricing page"), D("model_count", "Models priced", "number", None, "prices"),
                 D("min_input_per_mtok", "Min input price", "number", "USD / 1M tokens", "prices"), D("min_output_per_mtok", "Min output price", "number", "USD / 1M tokens", "prices"),
                 D("features", "Features", "list", None, "prices")],
    "hardware": [D("kind", "Kind"), D("manufacturer", "Manufacturer"), D("memory_gb", "Memory", "number", "GB"), D("memory_bandwidth_gbs", "Memory bandwidth", "number", "GB/s"),
                 D("tdp_watts", "TDP", "number", "W"), D("architecture", "Architecture"), D("memory_type", "Memory type"), D("runtimes", "Runtimes", "list"), D("release_date", "Release", "date")],
    "framework": [D("latest_version", "Latest version"), D("latest_release_at", "Latest release", "date"), D("license", "License"), D("language", "Language"),
                  D("metric.stars", "GitHub stars", "number"), D("metric.forks", "Forks", "number"), D("repository_url", "Repository")],
    "company": [D("country", "Country"), D("founded", "Founded", "date"), D("headquarters", "Headquarters"), D("org_kind", "Kind"), D("website", "Website"),
                D("model_count", "Models", "number", None, "graph"), D("paper_count", "Papers", "number", None, "graph")],
    "benchmark": [D("category", "Category"), D("metric", "Metric"), D("unit", "Unit"), D("task", "Task"), D("result_count", "Results", "number", None, "results")],
}


@router.get("")
@cached(300)
async def compare(request: Request, ids: str = Query(..., description="2–6 slugs or ids, comma-separated")) -> dict[str, Any]:
    keys = csv(ids)
    if not 2 <= len(keys) <= 6:
        raise ApiError(400, "ids must list between 2 and 6 entities")
    async with connection() as conn:
        rows = [await resolve_entity(conn, k) for k in keys]
        types = {r["entity_type"] for r in rows}
        etype = rows[0]["entity_type"]
        if etype in COMPANY_TYPES:
            etype = "company"
            types = {"company"}
        if len(types) != 1:
            raise ApiError(400, f"all entities must share one type, got {', '.join(sorted(types))}")
        dims = [dict(d) for d in DIMENSIONS.get(etype, [])]
        eids = [r["id"] for r in rows]
        items: list[dict[str, Any]] = []
        extra: dict[str, dict[str, Any]] = {eid: {} for eid in eids}
        prices_by: dict[str, list[dict[str, Any]]] = {eid: [] for eid in eids}
        results_by: dict[str, list[dict[str, Any]]] = {eid: [] for eid in eids}
        if etype == "model":
            for p in await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where p.model_id = any(cast(:ids as text[])) and p.valid_to is null order by p.input_per_mtok nulls last", ids=eids):
                prices_by[p["m_id"]].append(price_row(p))
            for eid, plist in prices_by.items():
                ins = [x["input_per_mtok"] for x in plist if x["input_per_mtok"] is not None]
                outs = [x["output_per_mtok"] for x in plist if x["output_per_mtok"] is not None]
                extra[eid] = {"best_input_per_mtok": min(ins) if ins else None, "best_output_per_mtok": min(outs) if outs else None,
                              "provider_count": len({x["provider"]["id"] for x in plist if x["provider"]})}
            res = await fetch_all(conn, f"select {RESULT_COLS} from {RESULT_FROM} where r.model_id = any(cast(:ids as text[])) and r.valid_to is null order by b.canonical_name, r.observed_at desc", ids=eids)
            per_bench: dict[str, dict[str, float]] = {}
            bench_meta: dict[str, dict[str, Any]] = {}
            for r in res:
                results_by[r["m_id"]].append(result_row(r))
                per_bench.setdefault(r["b_id"], {}).setdefault(r["m_id"], r["score"])
                bench_meta.setdefault(r["b_id"], {"slug": r["b_slug"], "name": r["b_canonical_name"], "unit": r.get("unit") or r.get("b_attributes", {}).get("unit"), "higher_is_better": r["higher_is_better"]})
            for bid, scores in per_bench.items():
                if all(eid in scores for eid in eids):
                    meta = bench_meta[bid]
                    dims.append({"key": f"bench:{meta['slug']}", "label": meta["name"], "kind": "number", "unit": meta["unit"], "source": "results", "higher_is_better": meta["higher_is_better"]})
                    for eid in eids:
                        extra[eid][f"bench:{meta['slug']}"] = scores[eid]
        elif etype == "provider":
            agg = await fetch_all(conn, """select provider_id, count(distinct model_id) as model_count, min(nullif(input_per_mtok, 0)) as min_input_per_mtok,
                                           min(nullif(output_per_mtok, 0)) as min_output_per_mtok, jsonb_agg(distinct k.key) filter (where k.key is not null) as features
                                           from prices p left join lateral jsonb_object_keys(p.features) k(key) on true
                                           where provider_id = any(cast(:ids as text[])) and valid_to is null group by 1""", ids=eids)
            for a in agg:
                extra[a["provider_id"]] = {"model_count": int(a["model_count"]), "min_input_per_mtok": a["min_input_per_mtok"], "min_output_per_mtok": a["min_output_per_mtok"], "features": a["features"] or []}
        elif etype == "company":
            agg = await fetch_all(conn, """select e.id, (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' and m.merged_into is null) as model_count,
                                           (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'paper' and m.merged_into is null) as paper_count
                                           from entities e where e.id = any(cast(:ids as text[]))""", ids=eids)
            for a in agg:
                extra[a["id"]] = {"model_count": int(a["model_count"]), "paper_count": int(a["paper_count"])}
        elif etype == "benchmark":
            agg = await fetch_all(conn, "select benchmark_id, count(*) as n from benchmark_results where benchmark_id = any(cast(:ids as text[])) and valid_to is null group by 1", ids=eids)
            for a in agg:
                extra[a["benchmark_id"]] = {"result_count": int(a["n"])}
        await enrich_provenance(conn, *(r.get("provenance") for r in rows))
        for r in rows:
            attrs = r.get("attributes") or {}
            prov = r.get("provenance") or {}
            values = {d["key"]: (extra[r["id"]].get(d["key"]) if d.get("source") != "attr" else attrs.get(d["key"])) for d in dims}
            item: dict[str, Any] = {"entity": entity_summary(r), "values": values, "provenance": {k: prov[k] for k in values if k in prov}}
            if etype == "model":
                item["prices"] = prices_by[r["id"]]
                item["results"] = results_by[r["id"]]
            items.append(item)
    return {"entity_type": etype, "dimensions": dims, "items": items}