"""Find-a-model (API 1.1): deterministic matching of canonical models against use-case criteria. No LLM, no composite "winner" score —
each match lists the observed facts (`why`) that satisfied each criterion and the observed dimensions; sorting is by number of satisfied
criteria, then by the best benchmark rank the model holds anywhere, then by release date.

Rules (documented in `RULES`, returned to the client as `filters_applied`):
    coding        current result on a coding-category benchmark, or capabilities/modalities mention code
    reasoning     attributes.reasoning is true, or a current result on a reasoning/math benchmark
    agentic       current result on an agentic benchmark, or tool_calling is true
    long_context  context_length ≥ context_min (default 128 000)
    vision        modalities (any direction) include image
    low_cost      cheapest current output price ≤ max_output_price, or in the bottom quartile of all cheapest output prices
    local         weights downloadable (open-weights / open-source / restricted-weights) and the estimated footprint fits memory_gb
    embeddings    modalities include embedding (or pipeline tag says so)
    chat          text modality (or unspecified) and reachable through ≥ 1 current deployment or downloadable weights
"""
from __future__ import annotations

from collections import defaultdict
from typing import Any

from sqlalchemy.ext.asyncio import AsyncConnection

from aiatlas.db import fetch_all
from aiatlas.ontology.licenses import LICENSES, normalize_license
from aiatlas.services import hardware_fit as hf
from aiatlas.services.frontier import all_primary_groups, rank_rows

USE_CASES = ("coding", "reasoning", "agentic", "long_context", "vision", "low_cost", "local", "embeddings", "chat")
DOWNLOADABLE = ("open-weights", "open-source", "restricted-weights", "restricted")
CODING_CATEGORIES = ("coding",)
REASONING_CATEGORIES = ("reasoning", "math")
AGENTIC_CATEGORIES = ("agentic",)
RULES = {
    "coding": "current result on a coding-category benchmark, or capabilities/modalities mention code",
    "reasoning": "attributes.reasoning = true, or a current result on a reasoning/math benchmark",
    "agentic": "current result on an agentic benchmark, or tool_calling = true",
    "long_context": "context_length ≥ context_min (default 128 000 tokens)",
    "vision": "modalities include image",
    "low_cost": "cheapest current output price ≤ max_output_price, or in the bottom quartile of all models' cheapest output prices",
    "local": "weights downloadable and the ESTIMATED footprint (services.hardware_fit) fits memory_gb at the requested quantisation",
    "embeddings": "modalities include embedding",
    "chat": "text modality and reachable (≥ 1 current deployment or downloadable weights)",
    "deployment": "local → downloadable weights · api → ≥ 1 current deployment · any → no constraint",
    "license": "commercial → licence key known in the ontology with commercial_use = true",
    "openness": "attributes.openness in the requested categories",
    "max_input_price / max_output_price": "cheapest current provider price (USD per 1M tokens) ≤ the bound",
    "modalities": "every requested modality present in modalities / modalities_input / modalities_output",
}


def _modalities(attrs: dict[str, Any]) -> set[str]:
    out: set[str] = set()
    for k in ("modalities", "modalities_input", "modalities_output"):
        v = attrs.get(k)
        if isinstance(v, list):
            out |= {str(x).strip().lower() for x in v}
        elif isinstance(v, str):
            out |= {x.strip().lower() for x in v.split(",") if x.strip()}
    if attrs.get("vision") is True:
        out.add("image")
    if attrs.get("audio") is True:
        out.add("audio")
    return {("image" if m in ("vision", "images") else "document" if m == "pdf" else m) for m in out}


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


def _mentions_code(attrs: dict[str, Any]) -> bool:
    caps = attrs.get("capabilities")
    text = " ".join(str(c) for c in caps).lower() if isinstance(caps, list) else str(caps or "").lower()
    return "code" in text or "code" in _modalities(attrs) or str(attrs.get("pipeline_tag") or "").lower().startswith("text-generation") and "coder" in str(attrs.get("family") or "").lower()


async def find_models(conn: AsyncConnection, *, use_case: str | None, deployment: str = "any", memory_gb: float | None = None, quant: str = "4bit",
                      context_min: int | None = None, license: str = "any", openness: list[str] | None = None, max_input_price: float | None = None,
                      max_output_price: float | None = None, modalities: list[str] | None = None, limit: int = 30) -> dict[str, Any]:
    models = await fetch_all(conn, """
        select e.id, e.slug, e.canonical_name, e.attributes, e.organization_id, o.slug as org_slug, o.canonical_name as org_name, e.first_seen_at
        from entities e left join entities o on o.id = e.organization_id where e.entity_type = 'model' and e.merged_into is null""")
    prices = await fetch_all(conn, """
        select p.model_id, min(p.input_per_mtok) filter (where p.input_per_mtok > 0) as min_input, min(p.output_per_mtok) filter (where p.output_per_mtok > 0) as min_output,
               count(distinct p.provider_id) as providers
        from prices p where p.valid_to is null group by 1""")
    price_by = {p["model_id"]: p for p in prices}
    groups = await all_primary_groups(conn)
    ranks: dict[str, dict[str, int]] = defaultdict(dict)       # model_id → {benchmark slug: rank}
    categories: dict[str, set[str]] = defaultdict(set)         # model_id → benchmark categories with a current result
    for g in groups.values():
        cat = (g["benchmark"].get("category") or "").lower()
        for r in rank_rows(g["rows"], g["higher_is_better"]):
            ranks[r["model_id"]][g["benchmark"]["slug"]] = r["rank"]
            if cat:
                categories[r["model_id"]].add(cat)
    outputs = sorted(float(p["min_output"]) for p in prices if p["min_output"] is not None)
    bottom_quartile = outputs[len(outputs) // 4] if outputs else None

    ctx_min = context_min or 128_000
    quant = hf.normalize_quant(quant)
    want_mods = {m.strip().lower() for m in (modalities or []) if m.strip()}
    open_set = set(openness or [])
    matches: list[dict[str, Any]] = []
    for m in models:
        attrs = m["attributes"] or {}
        why: list[str] = []
        failed = False
        mods = _modalities(attrs)
        pr = price_by.get(m["id"])
        cheapest_out = _num(pr["min_output"]) if pr else None
        cheapest_in = _num(pr["min_input"]) if pr else None
        providers = int(pr["providers"]) if pr else 0
        opn = str(attrs.get("openness") or "")
        downloadable = opn in DOWNLOADABLE
        ctx = _num(attrs.get("context_length"))
        cats = categories.get(m["id"], set())
        estimated_fit: dict[str, Any] | None = None

        # --- hard filters
        if deployment == "local" and not downloadable:
            continue
        if deployment == "api" and providers == 0:
            continue
        if open_set and opn not in open_set:
            continue
        if license == "commercial":
            key = attrs.get("license_key") or normalize_license(attrs.get("license"))
            info = LICENSES.get(key) if key else None
            if not info or info.commercial_use is not True:
                continue
            why.append(f"licence {info.label} allows commercial use")
        if max_input_price is not None:
            if cheapest_in is None or cheapest_in > max_input_price:
                continue
            why.append(f"cheapest input ${cheapest_in:g}/M ≤ ${max_input_price:g}")
        if max_output_price is not None and use_case != "low_cost":
            if cheapest_out is None or cheapest_out > max_output_price:
                continue
            why.append(f"cheapest output ${cheapest_out:g}/M ≤ ${max_output_price:g}")
        if want_mods and not want_mods <= mods:
            continue
        if want_mods:
            why.append("modalities include " + ", ".join(sorted(want_mods)))
        if memory_gb is not None and (use_case == "local" or deployment == "local"):
            estimated_fit = hf.fit_detailed(attrs, memory_gb, quant=quant, context=min(int(ctx or 8192), 8192))
            if estimated_fit is None or not estimated_fit["fits"]:
                continue
            why.append(f"estimated {estimated_fit['estimated_memory_gb']} GB at {quant} fits {memory_gb:g} GB (estimate)")

        # --- use-case rule (one criterion; `why` records the observed fact)
        if use_case == "coding":
            if cats & set(CODING_CATEGORIES):
                why.append("current result on a coding benchmark: " + ", ".join(sorted(b for b in ranks[m["id"]] if groups_cat(groups, b) == "coding")[:4]))
            elif _mentions_code(attrs):
                why.append("capabilities/modalities mention code")
            else:
                failed = True
        elif use_case == "reasoning":
            if attrs.get("reasoning") is True:
                why.append("attributes.reasoning = true")
            elif cats & set(REASONING_CATEGORIES):
                why.append("current result on a reasoning/math benchmark")
            else:
                failed = True
        elif use_case == "agentic":
            if cats & set(AGENTIC_CATEGORIES):
                why.append("current result on an agentic benchmark")
            elif attrs.get("tool_calling") is True:
                why.append("tool_calling = true")
            else:
                failed = True
        elif use_case == "long_context":
            if ctx is not None and ctx >= ctx_min:
                why.append(f"context_length {int(ctx):,} ≥ {ctx_min:,}")
            else:
                failed = True
        elif use_case == "vision":
            if "image" in mods:
                why.append("modalities include image")
            else:
                failed = True
        elif use_case == "low_cost":
            bound = max_output_price if max_output_price is not None else bottom_quartile
            if cheapest_out is not None and bound is not None and cheapest_out <= bound:
                why.append(f"cheapest output ${cheapest_out:g}/M ≤ ${bound:g}/M " + ("(bound)" if max_output_price is not None else "(bottom quartile of observed prices)"))
            else:
                failed = True
        elif use_case == "local":
            if not downloadable:
                failed = True
            else:
                why.append(f"openness {opn}: weights downloadable")
                if memory_gb is None:
                    estimated_fit = hf.fit_detailed(attrs, 64.0, quant=quant, context=8192)
                    if estimated_fit is not None:
                        why.append(f"estimated {estimated_fit['estimated_memory_gb']} GB at {quant} (reference 64 GB device, estimate)")
        elif use_case == "embeddings":
            if "embedding" in mods or "embedding" in str(attrs.get("pipeline_tag") or "").lower():
                why.append("modalities include embedding")
            else:
                failed = True
        elif use_case == "chat":
            if (not mods or "text" in mods) and (providers > 0 or downloadable):
                why.append("text modality and reachable (" + (f"{providers} providers" if providers else "downloadable weights") + ")")
            else:
                failed = True
        if failed:
            continue
        if context_min is not None and use_case != "long_context":
            if ctx is None or ctx < context_min:
                continue
            why.append(f"context_length {int(ctx):,} ≥ {context_min:,}")
        best_rank = min(ranks[m["id"]].values()) if ranks.get(m["id"]) else None
        matches.append({
            "model": {"id": m["id"], "slug": m["slug"], "name": m["canonical_name"], "entity_type": "model",
                      "organization": {"id": m["organization_id"], "slug": m["org_slug"], "name": m["org_name"]} if m["organization_id"] else None},
            "why": why,
            "observed": {"context_length": ctx, "parameter_count": _num(attrs.get("parameter_count")), "openness": opn or None, "license": attrs.get("license"),
                         "license_key": attrs.get("license_key") or normalize_license(attrs.get("license")), "modalities": sorted(mods), "reasoning": attrs.get("reasoning"),
                         "tool_calling": attrs.get("tool_calling"), "release_date": attrs.get("release_date"), "cheapest_input_per_mtok": cheapest_in,
                         "cheapest_output_per_mtok": cheapest_out, "providers": providers, "benchmark_ranks": dict(sorted(ranks.get(m["id"], {}).items())), "best_rank": best_rank},
            **({"estimated_fit": estimated_fit} if estimated_fit else {}),
            "_sort": (-len(why), best_rank if best_rank is not None else 10_000, -(_date_ord(attrs.get("release_date")))),
        })
    matches.sort(key=lambda x: x["_sort"])
    for x in matches:
        x.pop("_sort", None)
    return {"matches": matches[:limit], "total": len(matches),
            "filters_applied": {k: v for k, v in {"use_case": use_case, "deployment": deployment, "memory_gb": memory_gb, "quant": quant if memory_gb is not None or use_case == "local" else None,
                                                  "context_min": context_min, "license": license, "openness": openness or None, "max_input_price": max_input_price,
                                                  "max_output_price": max_output_price, "modalities": sorted(want_mods) or None}.items() if v not in (None, "any")},
            "rules": {k: RULES[k] for k in RULES if k == use_case or k not in USE_CASES},
            "note": "Deterministic filters over observed attributes, current prices and current benchmark results of canonical models; sorted by number of satisfied "
                    "criteria, then best benchmark rank, then release date. No composite score. Hardware fit is an estimate."}


def groups_cat(groups: dict[str, dict[str, Any]], slug: str) -> str:
    for g in groups.values():
        if g["benchmark"]["slug"] == slug:
            return (g["benchmark"].get("category") or "").lower()
    return ""


def _date_ord(v: Any) -> int:
    s = str(v or "")
    digits = "".join(ch for ch in s[:10] if ch.isdigit())
    return int(digits.ljust(8, "0")[:8]) if digits else 0


__all__ = ["RULES", "USE_CASES", "find_models"]