"""/models listing (canonical universe, filters on `entities.attributes`, facets), /models/{slug} alias, /models/{a}/diff/{b}.

API 1.1: the default universe is CANONICAL MODELS (`entity_type = 'model' and merged_into is null`). `include=artifacts` adds
checkpoints / quantisations / conversions as rows with `entity_type: 'artifact'` and a `canonical` summary."""
from __future__ import annotations

import asyncio
from typing import Any

from fastapi import APIRouter, Query, Request

from aiatlas.api.common import (
    ENTITY_COLS,
    ENTITY_FROM,
    MODEL_OR_ARTIFACT_UNIVERSE,
    MODEL_UNIVERSE,
    PAGINATION,
    ApiError,
    Pagination,
    attr_num,
    cached,
    entity_cols,
    entity_summary,
    flip_order,
    normalize,
    num_expr,
    openness_values,
    page,
)
from aiatlas.api.routers.entities import detail_for_type
from aiatlas.db import connection, fetch_all, fetch_val
from aiatlas.ontology.licenses import LICENSES, normalize_license
from aiatlas.services import cache

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

PARAMS = num_expr("e.attributes->>'parameter_count'")
CONTEXT = num_expr("e.attributes->>'context_length'")
DOWNLOADS = num_expr("e.attributes->>'metric.downloads'")
SORTS = {
    "updated": "e.updated_at desc", "name": "e.canonical_name asc", "params": f"{PARAMS} desc nulls last", "context": f"{CONTEXT} desc nulls last",
    "release": "e.attributes->>'release_date' desc nulls last", "quality": "coalesce((e.quality->>'score')::float, 0) desc", "downloads": f"{DOWNLOADS} desc nulls last",
    "first_seen": "e.first_seen_at desc",
    "cheapest": "(select min(p.output_per_mtok) from prices p where p.model_id = e.id and p.valid_to is null and p.output_per_mtok > 0) asc nulls last",
}
FAMILY_JOIN = "left join entities fam on fam.id = e.family_id"
CANONICAL_JOIN = "left join entities can on can.id = e.canonical_id left join entities cano on cano.id = can.organization_id"


def license_match_sql(param: str = "license") -> str:
    """Match a canonical licence key (`attributes.license_key`) OR any raw label the ontology maps to it OR a plain ilike on the raw label."""
    return f"(e.attributes->>'license_key' = :{param} or lower(e.attributes->>'license') = any(cast(:{param}_raw as text[])) or e.attributes->>'license' ilike :{param})"


def license_params(value: str) -> dict[str, Any]:
    key = value if value in LICENSES else (normalize_license(value) or value)
    info = LICENSES.get(key)
    raw = {key.lower(), value.lower()}
    if info:
        raw |= {a.lower() for a in info.aliases}
        if info.spdx:
            raw.add(info.spdx.lower())
    return {"license": key, "license_raw": sorted(raw)}


def model_filters(*, q: str | None, org: str | None, family: str | None, openness: str | None, modality: str | None, status: str | None,
                  min_params: float | None, max_params: float | None, min_context: int | None, year_from: int | None, year_to: int | None,
                  license: str | None, include_artifacts: bool = False, reasoning: bool | None = None, trust: str | None = None) -> tuple[list[str], dict[str, Any]]:
    where = [MODEL_OR_ARTIFACT_UNIVERSE if include_artifacts else MODEL_UNIVERSE]
    p: dict[str, Any] = {}
    if q:
        where.append("(e.canonical_name ilike :qlike or e.search @@ plainto_tsquery('simple', :q) or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")
        p["q"], p["qlike"] = q, f"%{q}%"
    if org:
        where.append("(eo.slug = :org or eo.id = :org or eo.canonical_name ilike :org)")
        p["org"] = org
    if family:
        # family slug / id (model_family entity via family_id) — falls back to the legacy `attributes.family` label
        where.append("(fam.slug = :family or fam.id = :family or fam.canonical_name ilike :family or e.attributes->>'family' ilike :family)")
        p["family"] = family
    if openness:
        where.append("e.attributes->>'openness' = any(cast(:openness as text[]))")
        p["openness"] = openness_values(openness)
    if modality:
        where.append("(e.attributes->'modalities' ? :modality or e.attributes->'modalities_input' ? :modality or e.attributes->'modalities_output' ? :modality)")
        p["modality"] = modality
    if status:
        where.append("e.status = any(cast(:status as text[]))")
        p["status"] = [v.strip() for v in status.split(",") if v.strip()]
    if min_params is not None:
        where.append(f"{PARAMS} >= :min_params")
        p["min_params"] = float(min_params)
    if max_params is not None:
        where.append(f"{PARAMS} <= :max_params")
        p["max_params"] = float(max_params)
    if min_context is not None:
        where.append(f"{CONTEXT} >= :min_context")
        p["min_context"] = float(min_context)
    if year_from is not None:
        where.append("left(e.attributes->>'release_date', 4) >= :yf")
        p["yf"] = str(year_from)
    if year_to is not None:
        where.append("left(e.attributes->>'release_date', 4) <= :yt")
        p["yt"] = str(year_to)
    if license:
        where.append(license_match_sql())
        p.update(license_params(license))
    if reasoning is not None:
        where.append("e.attributes->>'reasoning' = :reasoning")
        p["reasoning"] = "true" if reasoning else "false"
    if trust:
        where.append("e.identity_confidence = any(cast(:trust as text[]))")
        p["trust"] = [v.strip() for v in trust.split(",") if v.strip()]
    return where, p


def _canonical_licenses(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Fold raw licence labels into canonical keys (Apache 2.0 / apache-2.0 → Apache-2.0); unknown labels are kept raw under `raw: true`."""
    agg: dict[str, dict[str, Any]] = {}
    for r in rows:
        raw = r["value"]
        if raw in (None, ""):
            continue
        key = normalize_license(raw)
        k = key or str(raw)
        item = agg.setdefault(k, {"value": k, "label": LICENSES[key].label if key else str(raw), "category": LICENSES[key].category if key else "unknown",
                                  "count": 0, "raw_labels": [], **({} if key else {"raw": True})})
        item["count"] += int(r["count"])
        if str(raw) != k:
            item["raw_labels"].append(str(raw))
    return sorted(agg.values(), key=lambda x: (-x["count"], x["value"]))


async def model_facets(where_sql: str, params: dict[str, Any]) -> dict[str, Any]:
    async def run(sql: str) -> list[dict[str, Any]]:
        async with connection() as conn:
            return await fetch_all(conn, sql, **params)

    base = f"from {ENTITY_FROM} {FAMILY_JOIN} where {where_sql}"
    mods_from = (f"from {ENTITY_FROM} {FAMILY_JOIN} cross join lateral jsonb_array_elements_text(case when jsonb_typeof(e.attributes->'modalities') = 'array' "
                 f"then e.attributes->'modalities' else '[]'::jsonb end) m where {where_sql}")
    orgs, openness, mods, fams, years, lics, status, trust = await asyncio.gather(
        run(f"select eo.slug, eo.canonical_name as name, count(*) as count {base} and eo.id is not null group by 1, 2 order by 3 desc, 2 limit 60"),
        run(f"select e.attributes->>'openness' as value, count(*) as count {base} and e.attributes ? 'openness' group by 1 order by 2 desc"),
        run(f"select m.value, count(*) as count {mods_from} group by 1 order by 2 desc limit 30"),
        run(f"select coalesce(fam.slug, e.attributes->>'family') as value, coalesce(fam.canonical_name, e.attributes->>'family') as label, fam.id is not null as canonical, "
            f"count(*) as count {base} and (fam.id is not null or e.attributes ? 'family') group by 1, 2, 3 order by 4 desc, 2 limit 60"),
        run(f"select left(e.attributes->>'release_date', 4) as value, count(*) as count {base} and e.attributes ? 'release_date' group by 1 order by 1 desc limit 30"),
        run(f"select coalesce(e.attributes->>'license_key', e.attributes->>'license') as value, count(*) as count {base} and (e.attributes ? 'license' or e.attributes ? 'license_key') group by 1 order by 2 desc, 1 limit 80"),
        run(f"select e.status as value, count(*) as count {base} group by 1 order by 2 desc"),
        run(f"select e.identity_confidence as value, count(*) as count {base} group by 1 order by 2 desc"),
    )
    conv = lambda rows: [{"value": r["value"], "count": int(r["count"])} for r in rows if r["value"] not in (None, "")]
    return {"organizations": [{"slug": r["slug"], "name": r["name"], "count": int(r["count"])} for r in orgs], "openness": conv(openness), "modalities": conv(mods),
            "families": [{"value": r["value"], "label": r["label"], "canonical": bool(r["canonical"]), "count": int(r["count"])} for r in fams if r["value"]],
            "years": conv(years), "licenses": _canonical_licenses(lics), "status": conv(status),
            "trust": [{**x, "label": {"high": "Identity confirmed", "medium": "Identity probable", "low": "Identity uncertain"}.get(x["value"], x["value"])} for x in conv(trust)],
            "definitions": {"trust": "identity_confidence of the row: how sure AI Atlas is that this entry is one real model release (high | medium | low)",
                            "licenses": "canonical licence keys from the ontology; raw labels that could not be classified are flagged raw: true",
                            "families": "model_family entities (canonical: true) or legacy attribute labels (canonical: false)"}}


@router.get("")
@cached(300)
async def list_models(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), org: str | None = None, family: str | None = None,
                      openness: str | None = None, modality: str | None = None, status: str | None = None, min_params: float | None = Query(None, ge=0),
                      max_params: float | None = Query(None, ge=0), min_context: int | None = Query(None, ge=0), year_from: int | None = Query(None, ge=1950, le=2100),
                      year_to: int | None = Query(None, ge=1950, le=2100), license: str | None = None, sort: str = "updated", order: str = "",
                      facets: int = Query(0, ge=0, le=1), include: str | None = Query(None, description="`artifacts` restores the pre-1.1 universe (models + artifacts)"),
                      reasoning: int | None = Query(None, ge=0, le=1), trust: str | None = Query(None, description="identity_confidence: high,medium,low")) -> dict[str, Any]:
    if sort not in SORTS:
        raise ApiError(400, f"sort must be one of {', '.join(SORTS)}")
    include_artifacts = "artifacts" in {x.strip() for x in (include or "").split(",")}
    where, params = model_filters(q=q, org=org, family=family, openness=openness, modality=modality, status=status, min_params=min_params, max_params=max_params,
                                  min_context=min_context, year_from=year_from, year_to=year_to, license=license, include_artifacts=include_artifacts,
                                  reasoning=None if reasoning is None else bool(reasoning), trust=trust)
    order_sql = flip_order(SORTS[sort], order)
    where_sql = " and ".join(where)
    async with connection() as conn:
        rows = await fetch_all(conn, f"""select {ENTITY_COLS}, e.family_id, e.canonical_id, e.artifact_kind, e.identity_confidence,
                                                fam.slug as family_slug, fam.canonical_name as family_name, {entity_cols("can", "c_")},
                                                bp.input_per_mtok as bp_input, bp.output_per_mtok as bp_output, bp.provider_slug as bp_provider_slug, bp.provider_name as bp_provider_name,
                                                bp.provider_id as bp_provider_id, bp.providers as bp_providers
                                         from {ENTITY_FROM} {FAMILY_JOIN} {CANONICAL_JOIN}
                                         left join lateral (select p.input_per_mtok, p.output_per_mtok, p.provider_id, pv.slug as provider_slug, pv.canonical_name as provider_name,
                                                                   (select count(distinct q.provider_id) from prices q where q.model_id = e.id and q.valid_to is null) as providers
                                                            from prices p join entities pv on pv.id = p.provider_id
                                                            where p.model_id = e.id and p.valid_to is null and p.output_per_mtok > 0 order by p.output_per_mtok asc, p.input_per_mtok asc nulls last limit 1) bp on true
                                         where {where_sql} order by {order_sql}, e.id limit :lim offset :off""",
                               lim=p.limit, off=p.offset, **params)
        total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} {FAMILY_JOIN} where {where_sql}", **params)
    items = []
    for r in rows:
        s = entity_summary(r) or {}
        if r.get("family_id"):
            s["family"] = {"id": r["family_id"], "slug": r["family_slug"], "name": r["family_name"]}
        if r["entity_type"] == "artifact":
            s["artifact_kind"] = r.get("artifact_kind")
            s["canonical"] = entity_summary(r, "c_")
        s["identity_confidence"] = r.get("identity_confidence")
        s["best_price"] = ({"input_per_mtok": r["bp_input"], "output_per_mtok": r["bp_output"], "unit": "USD per 1M tokens",
                            "provider": {"id": r["bp_provider_id"], "slug": r["bp_provider_slug"], "name": r["bp_provider_name"]}, "providers": int(r["bp_providers"] or 0),
                            "note": "cheapest current offer by output price"} if r.get("bp_output") is not None else None)
        items.append(s)
    out = page(items, int(total or 0), p)
    out["universe"] = "models+artifacts" if include_artifacts else "canonical models"
    if facets:
        out["facets"] = await _facets_cached(where_sql, params)
    return out


async def _facets_cached(where_sql: str, params: dict[str, Any]) -> dict[str, Any]:
    key = "facets:models:v2:" + where_sql + ":" + repr(sorted(params.items()))
    hit = await cache.cache_get(key)
    if hit is not None:
        return hit
    value = normalize(await model_facets(where_sql, params))
    await cache.cache_set(key, value, 600)
    return value


@router.get("/{slug}")
@cached(300)
async def get_model(request: Request, slug: str) -> dict[str, Any]:
    """Accepts canonical models AND artifacts (an artifact detail carries `canonical` + `artifact_kind`); folded variants follow `merged_into`
    and report `redirected_from` so the web layer can 301."""
    return await detail_for_type(slug, ("model", "artifact"))


@router.get("/{a}/diff/{b}")
@cached(300)
async def model_diff(request: Request, a: str, b: str) -> dict[str, Any]:
    """Only the dimensions where two models differ, with a delta (numeric % change, list added/removed)."""
    from aiatlas.api.routers.compare import compare_entities

    body = await compare_entities([a, b], diff_only=True)
    items = body["items"]
    if len(items) != 2:
        raise ApiError(400, "two models are required")
    va, vb = items[0]["values"], items[1]["values"]
    dims = []
    for d in body["dimensions"]:
        x, y = va.get(d["key"]), vb.get(d["key"])
        dims.append({**d, "a": x, "b": y, "delta": _delta(x, y, d.get("kind"))})
    return {"a": items[0]["entity"], "b": items[1]["entity"], "dimensions": dims, "comparability": body.get("comparability"),
            "note": "Only dimensions with differing observed values; numeric delta = (b − a) / a; lists show added/removed elements."}


def _delta(x: Any, y: Any, kind: str | None) -> dict[str, Any] | None:
    if isinstance(x, list) or isinstance(y, list):
        sx, sy = {str(v) for v in (x or [])}, {str(v) for v in (y or [])}
        return {"added": sorted(sy - sx), "removed": sorted(sx - sy)}
    try:
        fx, fy = float(x), float(y)
    except (TypeError, ValueError):
        return None
    return {"absolute": fy - fx, "percent": round((fy - fx) / fx * 100, 2) if fx else None}


__all__ = ["CONTEXT", "PARAMS", "attr_num", "license_match_sql", "license_params", "model_filters", "router"]