"""Aggregate statistics. Only counts and shares leave this module, never rows.

Sessions are counted once (the newest row per session_id), flagged rows are left out,
and any bucket seen fewer than K_MIN times is folded into "other", so a rare
configuration can't be traced back to one visitor.
"""

from __future__ import annotations

import statistics
import threading
import time
from collections import Counter, defaultdict
from typing import Callable, Iterable

from .events import latest_sessions
from .recommend import EVAL_DATASET, PPL_BENCHMARK_VERSIONS, CommunitySignal

K_MIN = 5
BENCH_TYPES = ("mlx_benchmark", "mlx_benchmark_submission")


def _suppress(counter: Counter, k: int = K_MIN, top: int = 15) -> list[dict]:
    total = sum(counter.values())
    items, other = [], 0
    for key, n in counter.most_common():
        if key is None or n < k or len(items) >= top:
            other += n
        else:
            items.append({"key": key, "count": n, "share": round(n / total, 4) if total else 0})
    if other:
        items.append({"key": "other", "count": other, "share": round(other / total, 4) if total else 0})
    return items


def _present(sessions: list[dict], key: str) -> Counter:
    """Shares among visits that set the field (unset isn't a category)."""
    return Counter(s[key] for s in sessions if s.get(key) is not None)


def compute(rows: Iterable[dict], k: int = K_MIN) -> dict:
    rows = latest_sessions([r for r in rows if not r.get("suspicious_flags")])
    sessions = [r for r in rows if r.get("event_type") == "session"]
    by_type = Counter(r.get("event_type") for r in rows)

    quants = Counter(q for s in sessions for q in (s.get("quants_searched") or []))
    qtotal = sum(v for key, v in quants.items() if key != "unknown")
    share = lambda b: round(quants.get(b, 0) / qtotal, 4) if qtotal else None
    contexts = [s["target_context"] for s in sessions if s.get("target_context")]
    engaged = [s for s in sessions if s.get("models_viewed") or s.get("models_compared") or s.get("models_clicked")]
    picked = Counter(m for s in sessions for m in dict.fromkeys(
        [*(s.get("models_compared") or []), *(s.get("models_clicked") or [])]))
    hw = [s for s in sessions if s.get("gpu_capability_class")]

    return {
        "k_min": k,
        "totals": {
            "sessions": len(sessions),
            "engaged_sessions": len(engaged),
            "hardware_checks": len(hw),
            "browser_benchmarks": by_type.get("browser_benchmark", 0),
            "feedback": by_type.get("feedback", 0),
            "mlx_benchmarks": sum(by_type.get(t, 0) for t in BENCH_TYPES),
        },
        "quant_share": {b: share(b) for b in ("2-bit", "3-bit", "4-bit", "5-bit", "6-bit", "8-bit")},
        "average_target_context": round(statistics.mean(contexts)) if len(contexts) >= k else None,
        "families": _suppress(Counter(f for s in sessions for f in (s.get("families_searched") or [])), k),
        "quantizations": _suppress(quants, k),
        "sizes": _suppress(_present(sessions, "parameter_bucket"), k),
        "ram_classes": _suppress(_present(sessions, "hardware_memory_class"), k),
        "contexts": _suppress(_present(sessions, "target_context"), k),
        "priorities": _suppress(_present(sessions, "priority"), k),
        "selected_models": _suppress(picked, k),
        "gpu_capability": _suppress(Counter(s.get("gpu_capability_class") for s in hw), k),
        "webgpu_available": _suppress(Counter(s.get("webgpu_available") for s in hw), k),
    }


def community_signals(rows: Iterable[dict]) -> dict[str, CommunitySignal]:
    fb: dict[str, list[dict]] = defaultdict(list)
    bench: dict[str, list[float]] = defaultdict(list)
    ppl: dict[str, list[float]] = defaultdict(list)
    for r in rows:
        if r.get("suspicious_flags"):
            continue
        m = r.get("selected_model")
        if not m:
            continue
        if r.get("event_type") == "feedback" and r.get("tried") == "yes":
            fb[m].append(r)
        elif r.get("event_type") in BENCH_TYPES:
            if r.get("generation_tps"):
                bench[m].append(float(r["generation_tps"]))
            if (r.get("perplexity") and r.get("eval_dataset") == EVAL_DATASET
                    and r.get("benchmark_version") in PPL_BENCHMARK_VERSIONS):
                ppl[m].append(float(r["perplexity"]))
    out = {}
    for m in set(fb) | set(bench) | set(ppl):
        f = fb.get(m, [])
        rated = [x for x in f if x.get("outcome") in ("worked", "problem")]
        pos = sum(1 for x in rated if x.get("outcome") == "worked")
        b = bench.get(m, [])
        out[m] = CommunitySignal(
            feedback_count=len(f),
            positive_share=round(pos / len(rated), 3) if rated else None,
            benchmark_count=len(b),
            median_generation_tps=round(statistics.median(b), 1) if b else None,
            median_perplexity=round(statistics.median(ppl[m]), 4) if ppl.get(m) else None,
            perplexity_count=len(ppl.get(m, [])),
        )
    return out


class StatsCache:
    def __init__(self, rows_source: Callable[[], list[dict]], ttl: float = 60):
        self._src = rows_source
        self.ttl = ttl
        self._at = 0.0
        self._stats: dict = {}
        self._community: dict[str, CommunitySignal] = {}
        self._lock = threading.Lock()

    def _refresh(self):
        with self._lock:
            if time.time() - self._at < self.ttl and self._stats:
                return
            rows = self._src()
            self._stats = compute(rows)
            self._community = community_signals(rows)
            self._at = time.time()

    def invalidate(self):
        self._at = 0.0

    def stats(self) -> dict:
        self._refresh()
        return self._stats

    def community(self) -> dict[str, CommunitySignal]:
        self._refresh()
        return self._community