"""Recommendation scoring.

The engine only looks at measurable properties: estimated memory fit, bit width,
context, size, Hub popularity, recency and community data. It never looks at a
model's name, uploader or brand. `tests/test_recommend.py` enforces that.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Protocol

from .catalogue import ModelRecord
from .memory import MemoryEstimate, estimate
from .parsing import Arch, detect_arch, detect_quant, weights_key
from .model_detail import ModelDetail

PRIORITIES = ["balanced", "quality", "speed", "memory", "long_context"]

# How much each priority cares about each component. Rows sum to 1.
WEIGHTS: dict[str, dict[str, float]] = {
    "balanced":     {"fit": 0.30, "quant": 0.15, "context": 0.10, "capacity": 0.15, "popularity": 0.20, "recency": 0.05, "community": 0.05},
    "quality":      {"fit": 0.25, "quant": 0.25, "context": 0.05, "capacity": 0.25, "popularity": 0.12, "recency": 0.03, "community": 0.05},
    "speed":        {"fit": 0.35, "quant": 0.25, "context": 0.05, "capacity": 0.00, "popularity": 0.15, "recency": 0.05, "community": 0.15},
    "memory":       {"fit": 0.35, "quant": 0.35, "context": 0.05, "capacity": 0.00, "popularity": 0.15, "recency": 0.05, "community": 0.05},
    "long_context": {"fit": 0.25, "quant": 0.15, "context": 0.30, "capacity": 0.05, "popularity": 0.15, "recency": 0.05, "community": 0.05},
}

# Preference over nominal bit width per priority (1 = most preferred).
QUANT_PREF: dict[str, dict[int, float]] = {
    "quality":      {16: 1.0, 8: 0.95, 6: 0.85, 5: 0.75, 4: 0.6, 3: 0.35, 2: 0.15},
    "balanced":     {4: 1.0, 5: 0.95, 6: 0.85, 8: 0.75, 3: 0.55, 16: 0.45, 2: 0.25},
    "speed":        {4: 1.0, 3: 0.85, 5: 0.85, 6: 0.7, 2: 0.6, 8: 0.55, 16: 0.25},
    "memory":       {2: 1.0, 3: 0.9, 4: 0.8, 5: 0.6, 6: 0.5, 8: 0.3, 16: 0.1},
    "long_context": {4: 1.0, 5: 0.9, 6: 0.8, 3: 0.7, 8: 0.65, 2: 0.4, 16: 0.3},
}

def fit_score(ratio: float) -> float:
    """Continuous version of the fit classes, so 'a bit too big' beats 'far too big'."""
    if ratio < 0.70:
        return 1.0
    if ratio < 0.85:
        return 1.0 - (ratio - 0.70) / 0.15 * 0.2
    if ratio < 1.0:
        return 0.8 - (ratio - 0.85) / 0.15 * 0.4
    return max(0.0, 0.4 * (1 - (ratio - 1.0)))


@dataclass
class Query:
    family: str | None = None
    size: str | None = None
    quant: str | None = None
    context: int = 8192
    ram_gb: float | None = None
    ram_source: str | None = None  # confirmed | estimated | None
    priority: str = "balanced"


@dataclass
class CommunitySignal:
    feedback_count: int = 0
    positive_share: float | None = None
    benchmark_count: int = 0
    median_generation_tps: float | None = None
    median_perplexity: float | None = None  # on EVAL_DATASET
    perplexity_count: int = 0


EVAL_DATASET = "wikitext2-test-v1"
# Benchmark versions whose perplexity is comparable. mlxbench-1 left BOS off most windows,
# which inflated perplexity for models that have a BOS token (Gemma by 100x).
PPL_BENCHMARK_VERSIONS = ("mlxbench-2",)
EVAL_LABEL = "wikitext-2"
# Perplexity is only comparable between quantizations of the same base model (same tokenizer).
# Across bit widths the priority-based preference already applies, so measurements only move a
# model relative to measured peers of the SAME base model and bit width: each 1% lower perplexity
# than their mean adds this much to the quantization component, each 1% higher subtracts it.
# Unmeasured models are untouched, so measuring is neutral in expectation.
PPL_POINTS_PER_PCT = 0.02
PPL_MAX_SHIFT = 0.25


def perplexity_refs(community: dict, base_of: dict, bucket_of: dict | None = None) -> dict:
    """(base_model, quant bucket) -> [(model id, median perplexity)].

    Only quantizations of the same weights are comparable, so the base model is refined
    by ``weights_key`` of the repo name: a coder or abliterated finetune that declares the
    same base_model lands in its own group.
    """
    bucket_of = bucket_of or {}
    groups: dict[tuple, list] = {}
    for mid, sig in community.items():
        base = base_of.get(mid)
        if base and sig.median_perplexity:
            groups.setdefault((ppl_group(base, mid), bucket_of.get(mid)), []).append((mid, sig.median_perplexity))
    return groups


def ppl_group(base_model: str, model_id: str) -> str:
    return f"{base_model}::{weights_key(model_id)}"


@dataclass
class Scored:
    model: ModelRecord
    score: float
    components: dict[str, float]
    reasons: list[str]
    memory: MemoryEstimate
    detail_loaded: bool
    community: CommunitySignal = field(default_factory=CommunitySignal)
    quant: dict = field(default_factory=dict)  # effective quantization (config beats name/tags)

    def to_dict(self) -> dict:
        return {
            "model": self.model.public(),
            "score": round(self.score, 1),
            "components": {k: round(v, 3) for k, v in self.components.items()},
            "reasons": self.reasons,
            "memory": self.memory.to_dict(),
            "detail_loaded": self.detail_loaded,
            "quant": self.quant or self.model.quant,
            "community": self.community.__dict__,
        }


class RecommendationEngine(Protocol):
    version: str

    def score(self, models: list[ModelRecord], query: Query,
              details: dict[str, ModelDetail], community: dict[str, CommunitySignal]) -> list[Scored]: ...


def _months_since(iso: str | None, now: datetime) -> float | None:
    if not iso:
        return None
    try:
        dt = datetime.fromisoformat(iso.replace("Z", "+00:00"))
    except ValueError:
        return None
    return max(0.0, (now - dt).days / 30.4)


def _fmt_ctx(n: int) -> str:
    return f"{n // 1024}k" if n >= 1024 else str(n)


def _fmt_downloads(n: int) -> str:
    if n >= 1_000_000:
        return f"{n / 1e6:.1f}M"
    if n >= 1000:
        return f"{n / 1e3:.0f}k"
    return str(n)


class HeuristicEngine:
    version = "heuristic-v1"

    def __init__(self, now: datetime | None = None):
        self._now = now

    def memory_for(self, m: ModelRecord, query: Query, detail: ModelDetail | None) -> tuple[MemoryEstimate, dict]:
        quant = m.quant
        arch = Arch(None, None, 0, None, None, None, None, False)
        file_bytes = None
        if detail:
            arch = detect_arch(detail.config)
            file_bytes = detail.weight_bytes
            if detail.config:
                q = detect_quant(m.id, m.tags, detail.config)
                if q.bits is not None:
                    quant = q.to_dict()
        mem = estimate(params=m.params, bits=quant["bits"], mode=quant["mode"], file_bytes=file_bytes,
                       arch=arch, context=query.context, ram_gb=query.ram_gb)
        return mem, quant

    def norms(self, models, query, details=None) -> dict:
        """Set-wide reference values, computed once over ALL matches so a model's score
        doesn't depend on which page of results it is rescored in."""
        details = details or {}
        fitting = []
        for m in models:
            if not m.params or m.partial:
                continue
            if query.ram_gb is None:
                fitting.append(m.params)
                continue
            mem, _ = self.memory_for(m, query, details.get(m.id))
            if mem.fit in ("Comfortable", "Likely"):
                fitting.append(m.params)
        return {
            "max_downloads": max((m.downloads for m in models), default=0),
            "max_likes": max((m.likes for m in models), default=0),
            "max_fitting_params": max(fitting, default=0),
            "ppl_refs": {},
        }

    def score(self, models, query, details=None, community=None, norms=None) -> list[Scored]:
        details = details or {}
        community = community or {}
        now = self._now or datetime.now(timezone.utc)
        priority = query.priority if query.priority in WEIGHTS else "balanced"
        weights = WEIGHTS[priority]
        norms = norms or self.norms(models, query, details)
        max_dl, max_likes = norms["max_downloads"], norms["max_likes"]
        max_fit_params = norms["max_fitting_params"]
        ppl_refs = norms.get("ppl_refs") or {}
        out = []
        for m in models:
            detail = details.get(m.id)
            mem, quant = self.memory_for(m, query, detail)
            comm = community.get(m.id) or CommunitySignal()
            comp: dict[str, float] = {}
            reasons: list[str] = []

            # Hardware fit
            if mem.total_gb is None:
                comp["fit"] = 0.5
                reasons.append("Memory need unknown: some model information isn't available.")
            elif query.ram_gb is None:
                comp["fit"] = 0.6
                reasons.append(f"Needs roughly {mem.total_gb:.1f} GB at {_fmt_ctx(query.context)} context (estimate). "
                               "Set your memory to check fit.")
            else:
                comp["fit"] = fit_score(mem.ratio)
                src = "confirmed" if query.ram_source == "confirmed" else "estimated"
                reasons.append(f"{mem.fit} fit: ~{mem.total_gb:.1f} GB needed vs ~{mem.usable_gb:.0f} GB "
                               f"GPU-usable on a {query.ram_gb:g} GB Mac ({src}).")

            # Quantization vs priority
            bits = quant["bits"]
            if bits is None:
                comp["quant"] = 0.4
                reasons.append("Quantization unknown.")
            else:
                key = 16 if bits >= 16 else int(round(bits))
                comp["quant"] = QUANT_PREF[priority].get(key, 0.4)
                if query.ram_gb and mem.fit == "Unlikely":
                    comp["quant"] *= 0.25  # bit-width preference is moot if it won't load
                label = quant.get("label") or f"{key}-bit"
                if quant is not m.quant and m.quant.get("bits") and quant.get("bits") != m.quant.get("bits"):
                    reasons.append(f"Its config.json says {label}, although the name or tags say {m.quant['label']}.")
                note = {
                    "quality": "more bits keep more of the original model's quality",
                    "memory": "fewer bits mean a smaller download and less memory",
                    "speed": "fewer bits usually decode faster on unified memory, but speed depends on the model and Mac",
                    "long_context": "leaves more memory for the KV cache",
                    "balanced": "4 to 6 bits is the usual quality/memory middle ground",
                }[priority]
                reasons.append(f"{label}: {note}.")

            # Measured quality: compare with measured quantizations of the same base model and width
            ppl = comm.median_perplexity
            bucket = m.quant_bucket
            if ppl and m.base_model and bits is not None:
                group = ppl_group(m.base_model, m.id)
                peers = [p for mid, p in ppl_refs.get((group, bucket), []) if mid != m.id]
                if peers:
                    peer_mean = sum(peers) / len(peers)
                    delta_pct = (peer_mean - ppl) / peer_mean * 100  # positive = better than peers
                    shift = max(-PPL_MAX_SHIFT, min(PPL_MAX_SHIFT, PPL_POINTS_PER_PCT * delta_pct))
                    comp["quant"] = max(0.0, comp["quant"] + shift)
                    comp["measured_quality"] = round(shift, 4)
                    n = f"{len(peers)} other measured {bucket} quantization{'s' if len(peers) > 1 else ''}"
                    direction = "lower" if delta_pct >= 0 else "higher"
                    reasons.insert(1, f"Perplexity {ppl:.2f} on {EVAL_LABEL}: {abs(delta_pct):.1f}% {direction} than "
                                      f"the {n} of this base model (lower is better).")
                else:
                    others = [(mid, p) for (grp, _), items in ppl_refs.items() if grp == group
                              for mid, p in items if mid != m.id]
                    extra = f" Other widths of this base model measure {', '.join(f'{p:.2f}' for _, p in others[:3])}." if others else ""
                    reasons.append(f"Perplexity {ppl:.2f} on {EVAL_LABEL}; no other {bucket} quantization of this base "
                                   f"model has been measured to compare against.{extra}")

            # Context
            if mem.exceeds_model_context:
                comp["context"] = 0.0
                reasons.append(f"The model's config lists a shorter maximum context than {_fmt_ctx(query.context)}.")
            elif mem.total_gb is not None and mem.kv_gb is not None:
                if query.ram_gb:
                    headroom = 1 - mem.total_gb / max(mem.usable_gb or 1, 0.1)
                    comp["context"] = max(0.0, min(1.0, 0.5 + headroom))
                else:
                    share = mem.kv_gb / max(mem.total_gb, 0.01)
                    comp["context"] = max(0.0, 1 - share)
                bound = " (rough, architecture unknown)" if mem.kv_rough else ""
                reasons.append(f"KV cache at {_fmt_ctx(query.context)}: ~{mem.kv_gb:.1f} GB{bound}.")
            else:
                comp["context"] = 0.5

            # Capacity: among models that fit, more parameters usually means a more capable model
            if query.ram_gb is None:
                comp["capacity"] = 0.5  # without a memory budget, "bigger" would just mean "won't fit"
            elif m.params and max_fit_params and mem.fit not in ("Borderline", "Unlikely"):
                comp["capacity"] = min(1.0, math.log10(1 + m.params / 1e8) / math.log10(1 + max_fit_params / 1e8))
                if weights["capacity"] and m.params < 0.5 * max_fit_params:
                    reasons.append(f"{m.params / 1e9:.1f}B parameters; models up to {max_fit_params / 1e9:.0f}B "
                                   f"in these results fit your Mac, and larger models are usually more capable.")
            elif m.params and max_fit_params:
                comp["capacity"] = 0.0
            else:
                comp["capacity"] = 0.5

            # Popularity: downloads and likes, log-scaled against the best in this result set
            dl = math.log10(m.downloads + 1) / math.log10(max_dl + 1) if max_dl > 0 else 0.0
            lk = math.log10(m.likes + 1) / math.log10(max_likes + 1) if max_likes > 0 else 0.0
            comp["popularity"] = 0.7 * dl + 0.3 * lk
            if m.downloads:
                reasons.append(f"{_fmt_downloads(m.downloads)} Hugging Face downloads.")

            # Recency: half-life of a year
            months = _months_since(m.modified, now)
            comp["recency"] = 0.5 ** (months / 12) if months is not None else 0.3

            # Community data (zero until people contribute)
            if comm.benchmark_count or comm.feedback_count:
                c = min(1.0, math.log10(1 + comm.benchmark_count + comm.feedback_count) / 2)
                if comm.positive_share is not None:
                    c = 0.5 * c + 0.5 * comm.positive_share
                comp["community"] = c
                bits_ = []
                if comm.benchmark_count:
                    tps = f", median {comm.median_generation_tps:.0f} tok/s" if comm.median_generation_tps else ""
                    bits_.append(f"{comm.benchmark_count} MLX benchmark(s){tps}")
                if comm.feedback_count:
                    bits_.append(f"{comm.feedback_count} community report(s)")
                reasons.append("Community: " + "; ".join(bits_) + ".")
            else:
                comp["community"] = 0.0

            score = min(100.0, 100 * sum(weights[k] * comp[k] for k in weights))
            if m.partial:
                score *= 0.6
                reasons.insert(0, "This repo holds far fewer weights than its name suggests. It may be an add-on "
                                  "(draft head, adapter or partial upload) rather than a complete model.")
            if query.ram_gb and mem.fit == "Unlikely":
                score *= 0.5  # never recommend something that likely won't load over something that will
            out.append(Scored(m, score, comp, reasons, mem, detail is not None and detail.config is not None, comm, quant))
        return out


def sort_results(scored: list[Scored], sort: str) -> list[Scored]:
    if sort == "popular":
        key = lambda s: (-s.model.downloads, s.model.id)
    elif sort == "recent":
        key = lambda s: (-(_ts(s.model.modified)), s.model.id)
    elif sort == "community":
        key = lambda s: (-(s.community.benchmark_count + s.community.feedback_count), -s.score, s.model.id)
    elif sort == "all":
        key = lambda s: s.model.id.lower()
    else:
        key = lambda s: (-s.score, -s.model.downloads, s.model.id)
    return sorted(scored, key=key)


def _ts(iso: str | None) -> float:
    if not iso:
        return 0.0
    try:
        return datetime.fromisoformat(iso.replace("Z", "+00:00")).timestamp()
    except ValueError:
        return 0.0