"""Telemetry-driven adaptive defaults for OBLITERATUS.

Fetches community telemetry from the HuggingFace Hub dataset and analyzes
historical runs to recommend the best abliteration method and hyperparameters
for a given model architecture.

Architecture bucketing:
  Records are grouped by (arch_class, reasoning_class, param_bucket) where
  param_bucket is a coarse size tier (tiny/small/medium/large/frontier).
  Within each bucket, methods are ranked by composite score and the
  best-performing hyperparameter ranges are extracted.

The ``get_adaptive_recommendation()`` function returns an
``AdaptiveRecommendation`` that the pipeline/UI can apply on top of
(or instead of) the static research-grounded defaults in
``architecture_profiles.py``.

Data flow:
  HF Hub (OBLITERATUS-TELEMETRY)  ──►  fetch_hub_records()
       │                                      │
       ▼                                      ▼
  Local JSONL cache  ──────────►  build_knowledge_base()
                                         │
                                         ▼
                              get_adaptive_recommendation()
                                         │
                                         ▼
                              AdaptiveRecommendation
                              (best method, overrides, confidence)
"""

from __future__ import annotations

import logging
import statistics
import time
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)

# ── Cache config ──────────────────────────────────────────────────────────

_CACHE_TTL_S = 600  # 10 minutes — telemetry doesn't change that fast
_cache: dict[str, Any] = {}
_cache_ts: float = 0.0

# Minimum records per bucket to trust the recommendation
_MIN_RECORDS_FOR_CONFIDENCE = 5
_HIGH_CONFIDENCE_RECORDS = 20


# ── Size bucketing ────────────────────────────────────────────────────────

def _param_bucket(total_params_b: float) -> str:
    """Coarse size tier matching presets.py tiers."""
    if total_params_b <= 0.5:
        return "tiny"
    if total_params_b <= 4:
        return "small"
    if total_params_b <= 16:
        return "medium"
    if total_params_b <= 80:
        return "large"
    return "frontier"


def _extract_arch_key(record: dict) -> tuple[str, str, str] | None:
    """Extract (arch_class, reasoning_class, param_bucket) from a telemetry record.

    Returns None if the record lacks enough information to classify.
    """
    model = record.get("model", {})
    if isinstance(model, str):
        # Schema v1 — just model name, can't reliably bucket
        return None

    arch_str = model.get("architecture", "")
    num_layers = model.get("num_layers", 0)
    hidden_size = model.get("hidden_size", 0)
    total_params = model.get("total_params", 0)

    # Estimate params in billions
    if total_params > 0:
        params_b = total_params / 1e9
    elif num_layers > 0 and hidden_size > 0:
        # Rough estimate: 12 * hidden² * num_layers (transformer scaling)
        params_b = (12 * hidden_size**2 * num_layers) / 1e9
    else:
        return None

    # Detect architecture class from the architecture string or model config
    arch_lower = arch_str.lower()
    moe_keywords = {"moe", "mixtral", "qwen2_moe", "qwen3_moe", "deepseek_v2",
                    "deepseek_v3", "dbrx", "grok", "jamba", "arctic", "olmoe",
                    "switch", "llama4"}
    is_moe = any(kw in arch_lower for kw in moe_keywords)

    # Check method_config for per_expert_directions as MoE signal
    mc = record.get("method_config", {})
    if mc.get("per_expert_directions"):
        is_moe = True

    if is_moe:
        arch_class = "large_moe" if params_b > 100 else "small_moe"
    else:
        arch_class = "dense"

    # Detect reasoning from analysis insights or architecture name
    analysis = record.get("analysis_insights", {})
    reasoning_class = "standard"
    reasoning_keywords = {"reason", "think", "cot", "r1", "qwq", "o1", "o3"}
    if any(kw in arch_lower for kw in reasoning_keywords):
        reasoning_class = "reasoning"
    if analysis.get("cot_aware") or mc.get("cot_aware"):
        reasoning_class = "reasoning"

    return (arch_class, reasoning_class, _param_bucket(params_b))


# ── Composite scoring (same as tourney.py) ────────────────────────────────

def _composite_score(qm: dict[str, Any]) -> float:
    """Score a run on [0, 1].  Higher is better."""
    rr = qm.get("refusal_rate")
    co = qm.get("coherence")
    kl = qm.get("kl_divergence")
    pp = qm.get("perplexity")

    refusal_score = (1.0 - rr) if rr is not None else 0.0
    coherence_score = co if co is not None else 0.0
    kl_score = 1.0 / (1.0 + kl) if kl is not None else 0.5
    ppl_score = 1.0 / (1.0 + pp / 100.0) if pp is not None else 0.5

    return (
        refusal_score * 0.4
        + coherence_score * 0.3
        + kl_score * 0.2
        + ppl_score * 0.1
    )


# ── Data structures ──────────────────────────────────────────────────────


@dataclass
class MethodStats:
    """Aggregated statistics for one method within an architecture bucket."""

    method: str
    n_runs: int = 0
    scores: list[float] = field(default_factory=list)
    refusal_rates: list[float] = field(default_factory=list)
    coherences: list[float] = field(default_factory=list)
    kl_divergences: list[float] = field(default_factory=list)
    perplexities: list[float] = field(default_factory=list)
    configs: list[dict[str, Any]] = field(default_factory=list)

    @property
    def mean_score(self) -> float:
        return statistics.mean(self.scores) if self.scores else 0.0

    @property
    def best_score(self) -> float:
        return max(self.scores) if self.scores else 0.0

    @property
    def median_score(self) -> float:
        return statistics.median(self.scores) if self.scores else 0.0

    def best_config_ranges(self) -> dict[str, Any]:
        """Extract the hyperparameter ranges from top-performing runs.

        Takes the top 25% of runs by composite score and returns the median
        value for each numeric config key, or the mode for booleans.
        """
        if not self.configs or not self.scores:
            return {}

        # Pair scores with configs and take top 25%
        paired = sorted(zip(self.scores, self.configs), key=lambda x: x[0], reverse=True)
        top_n = max(1, len(paired) // 4)
        top_configs = [c for _, c in paired[:top_n]]

        ranges: dict[str, Any] = {}
        all_keys = set()
        for c in top_configs:
            all_keys.update(c.keys())

        for key in all_keys:
            values = [c[key] for c in top_configs if key in c and c[key] is not None]
            if not values:
                continue

            if all(isinstance(v, bool) for v in values):
                # Mode for booleans
                true_count = sum(1 for v in values if v)
                ranges[key] = true_count > len(values) / 2
            elif all(isinstance(v, (int, float)) for v in values):
                # Median for numerics
                ranges[key] = statistics.median(values)
                # Round ints back to ints
                if all(isinstance(v, int) for v in values):
                    ranges[key] = int(round(ranges[key]))
            # Skip strings and other types

        return ranges


@dataclass
class BucketKnowledge:
    """Everything we know about one architecture bucket from telemetry."""

    arch_key: tuple[str, str, str]  # (arch_class, reasoning_class, param_bucket)
    methods: dict[str, MethodStats] = field(default_factory=dict)
    total_runs: int = 0

    @property
    def best_method(self) -> str | None:
        """Method with highest mean composite score (min 3 runs)."""
        candidates = [
            (name, ms) for name, ms in self.methods.items()
            if ms.n_runs >= 3
        ]
        if not candidates:
            # Fall back to any method with runs
            candidates = [(name, ms) for name, ms in self.methods.items() if ms.n_runs > 0]
        if not candidates:
            return None
        return max(candidates, key=lambda x: x[1].mean_score)[0]

    @property
    def ranked_methods(self) -> list[tuple[str, MethodStats]]:
        """All methods ranked by mean score, descending."""
        return sorted(
            self.methods.items(),
            key=lambda x: x[1].mean_score,
            reverse=True,
        )


@dataclass
class AdaptiveRecommendation:
    """A telemetry-driven recommendation for a specific model."""

    # What we recommend
    recommended_method: str
    method_overrides: dict[str, Any]

    # How confident we are
    confidence: str           # "high", "medium", "low", "none"
    n_records: int            # total records in bucket
    n_method_records: int     # records for this specific method

    # Context
    arch_key: tuple[str, str, str]
    bucket_label: str         # human-readable e.g. "Dense Standard Medium"
    method_ranking: list[tuple[str, float]]  # [(method, mean_score), ...]

    # Best metrics seen in this bucket
    best_refusal_rate: float | None = None
    best_coherence: float | None = None

    # Explanation
    reason: str = ""

    def to_dict(self) -> dict:
        return {
            "recommended_method": self.recommended_method,
            "method_overrides": self.method_overrides,
            "confidence": self.confidence,
            "n_records": self.n_records,
            "n_method_records": self.n_method_records,
            "arch_key": list(self.arch_key),
            "bucket_label": self.bucket_label,
            "method_ranking": self.method_ranking,
            "best_refusal_rate": self.best_refusal_rate,
            "best_coherence": self.best_coherence,
            "reason": self.reason,
        }


# ── Knowledge base construction ──────────────────────────────────────────


def build_knowledge_base(
    records: list[dict[str, Any]] | None = None,
) -> dict[tuple[str, str, str], BucketKnowledge]:
    """Build per-bucket knowledge from telemetry records.

    If *records* is None, fetches from local + Hub automatically.
    """
    if records is None:
        records = _fetch_all_records()

    buckets: dict[tuple[str, str, str], BucketKnowledge] = {}

    for record in records:
        # Skip errored runs
        if record.get("error"):
            continue

        arch_key = _extract_arch_key(record)
        if arch_key is None:
            continue

        method = record.get("method", "")
        if not method:
            continue

        qm = record.get("quality_metrics", {})
        if not qm:
            continue

        score = _composite_score(qm)

        if arch_key not in buckets:
            buckets[arch_key] = BucketKnowledge(arch_key=arch_key)

        bucket = buckets[arch_key]
        bucket.total_runs += 1

        if method not in bucket.methods:
            bucket.methods[method] = MethodStats(method=method)

        ms = bucket.methods[method]
        ms.n_runs += 1
        ms.scores.append(score)

        rr = qm.get("refusal_rate")
        if rr is not None:
            ms.refusal_rates.append(rr)
        co = qm.get("coherence")
        if co is not None:
            ms.coherences.append(co)
        kl = qm.get("kl_divergence")
        if kl is not None:
            ms.kl_divergences.append(kl)
        pp = qm.get("perplexity")
        if pp is not None:
            ms.perplexities.append(pp)

        mc = record.get("method_config", {})
        if mc:
            ms.configs.append(mc)

    return buckets


def _fetch_all_records() -> list[dict[str, Any]]:
    """Fetch telemetry from local file + Hub, with caching."""
    global _cache, _cache_ts

    now = time.time()
    if _cache.get("records") is not None and (now - _cache_ts) < _CACHE_TTL_S:
        return _cache["records"]

    records: list[dict[str, Any]] = []

    # Local records
    try:
        from obliteratus.telemetry import read_telemetry
        records.extend(read_telemetry())
    except Exception as e:
        logger.debug("Failed to read local telemetry: %s", e)

    # Hub records
    try:
        from obliteratus.telemetry import fetch_hub_records
        hub = fetch_hub_records()
        records.extend(hub)
    except Exception as e:
        logger.debug("Failed to fetch Hub telemetry: %s", e)

    # Deduplicate by (session_id, timestamp)
    seen: set[tuple[str, str]] = set()
    deduped = []
    for r in records:
        key = (r.get("session_id", ""), r.get("timestamp", ""))
        if key not in seen:
            seen.add(key)
            deduped.append(r)

    _cache["records"] = deduped
    _cache_ts = now
    return deduped


# ── Recommendation engine ────────────────────────────────────────────────


def get_adaptive_recommendation(
    arch_class: str,
    reasoning_class: str,
    total_params_b: float,
    model_name: str = "",
    knowledge: dict[tuple[str, str, str], BucketKnowledge] | None = None,
) -> AdaptiveRecommendation:
    """Get a telemetry-based recommendation for the given architecture.

    Looks up the closest bucket in the knowledge base and returns the
    best-performing method + hyperparameter overrides.

    Falls through to broader buckets if the exact match has too few records:
      1. Exact match: (arch_class, reasoning_class, param_bucket)
      2. Size-agnostic: (arch_class, reasoning_class, "*")
      3. Arch-only: (arch_class, "*", "*")

    Args:
        arch_class: "dense", "small_moe", or "large_moe"
        reasoning_class: "standard" or "reasoning"
        total_params_b: Total params in billions
        model_name: Optional, for model-specific matching
        knowledge: Pre-built knowledge base (fetches if None)
    """
    if knowledge is None:
        knowledge = build_knowledge_base()

    param_bucket = _param_bucket(total_params_b)
    bucket_label = f"{arch_class.replace('_', ' ').title()} {reasoning_class.title()} {param_bucket.title()}"

    # Try exact match first, then broaden
    candidates = [
        (arch_class, reasoning_class, param_bucket),
    ]

    bucket = None
    used_key = None
    for key in candidates:
        if key in knowledge and knowledge[key].total_runs >= _MIN_RECORDS_FOR_CONFIDENCE:
            bucket = knowledge[key]
            used_key = key
            break

    # Fall back: merge all buckets that share (arch_class, reasoning_class)
    if bucket is None:
        merged = BucketKnowledge(arch_key=(arch_class, reasoning_class, "*"))
        for key, bkt in knowledge.items():
            if key[0] == arch_class and key[1] == reasoning_class:
                for method_name, ms in bkt.methods.items():
                    if method_name not in merged.methods:
                        merged.methods[method_name] = MethodStats(method=method_name)
                    target = merged.methods[method_name]
                    target.n_runs += ms.n_runs
                    target.scores.extend(ms.scores)
                    target.refusal_rates.extend(ms.refusal_rates)
                    target.coherences.extend(ms.coherences)
                    target.kl_divergences.extend(ms.kl_divergences)
                    target.perplexities.extend(ms.perplexities)
                    target.configs.extend(ms.configs)
                merged.total_runs += bkt.total_runs
        if merged.total_runs >= _MIN_RECORDS_FOR_CONFIDENCE:
            bucket = merged
            used_key = merged.arch_key
            bucket_label = f"{arch_class.replace('_', ' ').title()} {reasoning_class.title()} (all sizes)"

    # Last resort: merge all buckets that share arch_class
    if bucket is None:
        merged = BucketKnowledge(arch_key=(arch_class, "*", "*"))
        for key, bkt in knowledge.items():
            if key[0] == arch_class:
                for method_name, ms in bkt.methods.items():
                    if method_name not in merged.methods:
                        merged.methods[method_name] = MethodStats(method=method_name)
                    target = merged.methods[method_name]
                    target.n_runs += ms.n_runs
                    target.scores.extend(ms.scores)
                    target.refusal_rates.extend(ms.refusal_rates)
                    target.coherences.extend(ms.coherences)
                    target.kl_divergences.extend(ms.kl_divergences)
                    target.perplexities.extend(ms.perplexities)
                    target.configs.extend(ms.configs)
                merged.total_runs += bkt.total_runs
        if merged.total_runs > 0:
            bucket = merged
            used_key = merged.arch_key
            bucket_label = f"{arch_class.replace('_', ' ').title()} (all)"

    # No data at all
    if bucket is None or not bucket.methods:
        return AdaptiveRecommendation(
            recommended_method="",
            method_overrides={},
            confidence="none",
            n_records=0,
            n_method_records=0,
            arch_key=(arch_class, reasoning_class, param_bucket),
            bucket_label=bucket_label,
            method_ranking=[],
            reason="No telemetry data available for this architecture.",
        )

    # Get best method
    best_method = bucket.best_method
    if not best_method:
        return AdaptiveRecommendation(
            recommended_method="",
            method_overrides={},
            confidence="none",
            n_records=bucket.total_runs,
            n_method_records=0,
            arch_key=used_key or (arch_class, reasoning_class, param_bucket),
            bucket_label=bucket_label,
            method_ranking=[],
            reason="Telemetry records found but no method has enough runs.",
        )

    ms = bucket.methods[best_method]

    # Extract best hyperparams from top runs
    overrides = ms.best_config_ranges()

    # Confidence level
    if ms.n_runs >= _HIGH_CONFIDENCE_RECORDS:
        confidence = "high"
    elif ms.n_runs >= _MIN_RECORDS_FOR_CONFIDENCE:
        confidence = "medium"
    else:
        confidence = "low"

    # Method ranking
    ranking = [
        (name, stats.mean_score)
        for name, stats in bucket.ranked_methods
    ]

    # Best metrics seen
    best_rr = min(ms.refusal_rates) if ms.refusal_rates else None
    best_co = max(ms.coherences) if ms.coherences else None

    # Build explanation
    runner_up = ranking[1] if len(ranking) > 1 else None
    reason_parts = [
        f"Based on {bucket.total_runs} community runs for {bucket_label}.",
        f"`{best_method}` achieves a mean composite score of {ms.mean_score:.4f} "
        f"across {ms.n_runs} runs.",
    ]
    if runner_up:
        reason_parts.append(
            f"Runner-up: `{runner_up[0]}` ({runner_up[1]:.4f})."
        )
    if best_rr is not None:
        reason_parts.append(f"Best refusal rate seen: {best_rr:.1%}.")
    if overrides:
        override_strs = [f"{k}={v}" for k, v in sorted(overrides.items())]
        reason_parts.append(f"Optimal hyperparams from top runs: {', '.join(override_strs[:6])}")

    return AdaptiveRecommendation(
        recommended_method=best_method,
        method_overrides=overrides,
        confidence=confidence,
        n_records=bucket.total_runs,
        n_method_records=ms.n_runs,
        arch_key=used_key or (arch_class, reasoning_class, param_bucket),
        bucket_label=bucket_label,
        method_ranking=ranking,
        best_refusal_rate=best_rr,
        best_coherence=best_co,
        reason=" ".join(reason_parts),
    )


# ── Cross-architecture insights ──────────────────────────────────────────


def get_global_insights(
    knowledge: dict[tuple[str, str, str], BucketKnowledge] | None = None,
) -> dict[str, Any]:
    """Compute cross-architecture insights from all telemetry.

    Returns a summary dict with:
      - overall_best_methods: top methods across all architectures
      - architecture_breakdown: per-bucket summaries
      - total_records: total telemetry records analyzed
      - hyperparameter_trends: keys that consistently appear in top configs
    """
    if knowledge is None:
        knowledge = build_knowledge_base()

    total_records = sum(b.total_runs for b in knowledge.values())

    # Global method scores (weighted by bucket size)
    global_method_scores: dict[str, list[float]] = {}
    for bucket in knowledge.values():
        for name, ms in bucket.methods.items():
            if name not in global_method_scores:
                global_method_scores[name] = []
            global_method_scores[name].extend(ms.scores)

    overall_ranking = sorted(
        [
            (name, statistics.mean(scores), len(scores))
            for name, scores in global_method_scores.items()
            if scores
        ],
        key=lambda x: x[1],
        reverse=True,
    )

    # Per-bucket summaries
    arch_breakdown = {}
    for key, bucket in sorted(knowledge.items()):
        label = f"{key[0]} / {key[1]} / {key[2]}"
        best = bucket.best_method
        arch_breakdown[label] = {
            "total_runs": bucket.total_runs,
            "best_method": best,
            "best_score": bucket.methods[best].mean_score if best and best in bucket.methods else 0,
            "n_methods_tested": len(bucket.methods),
        }

    # Hyperparameter trends across top runs
    all_top_configs: list[dict] = []
    for bucket in knowledge.values():
        for ms in bucket.methods.values():
            if ms.configs and ms.scores:
                paired = sorted(zip(ms.scores, ms.configs), key=lambda x: x[0], reverse=True)
                top_n = max(1, len(paired) // 4)
                all_top_configs.extend(c for _, c in paired[:top_n])

    hp_trends: dict[str, Any] = {}
    if all_top_configs:
        all_keys = set()
        for c in all_top_configs:
            all_keys.update(c.keys())
        for key in sorted(all_keys):
            values = [c[key] for c in all_top_configs if key in c and c[key] is not None]
            if not values:
                continue
            if all(isinstance(v, bool) for v in values):
                true_pct = sum(1 for v in values if v) / len(values)
                hp_trends[key] = {"type": "bool", "true_pct": round(true_pct, 2), "n": len(values)}
            elif all(isinstance(v, (int, float)) for v in values):
                hp_trends[key] = {
                    "type": "numeric",
                    "median": round(statistics.median(values), 4),
                    "mean": round(statistics.mean(values), 4),
                    "min": min(values),
                    "max": max(values),
                    "n": len(values),
                }

    return {
        "total_records": total_records,
        "overall_best_methods": [
            {"method": name, "mean_score": round(score, 4), "n_runs": n}
            for name, score, n in overall_ranking
        ],
        "architecture_breakdown": arch_breakdown,
        "hyperparameter_trends": hp_trends,
    }


# ── Format helpers ────────────────────────────────────────────────────────


def format_recommendation(rec: AdaptiveRecommendation) -> str:
    """Format a recommendation as a human-readable markdown string."""
    if rec.confidence == "none":
        return (
            f"**No telemetry data** for {rec.bucket_label}.\n\n"
            "Using research-grounded defaults from `architecture_profiles.py`.\n"
            "Run some abliterations and the adaptive system will learn!"
        )

    confidence_emoji = {"high": "HIGH", "medium": "MEDIUM", "low": "LOW"}
    conf = confidence_emoji.get(rec.confidence, rec.confidence.upper())

    lines = [
        f"### Adaptive Recommendation [{conf} confidence]",
        f"**Architecture bucket:** {rec.bucket_label}",
        f"**Based on:** {rec.n_records} community runs",
        "",
        f"**Recommended method:** `{rec.recommended_method}` "
        f"(score: {rec.method_ranking[0][1]:.4f}, {rec.n_method_records} runs)",
        "",
    ]

    if len(rec.method_ranking) > 1:
        lines.append("**Method ranking:**")
        lines.append("| Rank | Method | Mean Score | Runs |")
        lines.append("|------|--------|------------|------|")
        for i, (name, score) in enumerate(rec.method_ranking[:8], 1):
            # Get run count from the knowledge (not stored directly, but we have n_method_records for winner)
            lines.append(f"| {i} | `{name}` | {score:.4f} | — |")
        lines.append("")

    if rec.method_overrides:
        lines.append("**Optimal hyperparameters** (from top 25% of runs):")
        for k, v in sorted(rec.method_overrides.items()):
            lines.append(f"  - `{k}`: {v}")
        lines.append("")

    if rec.best_refusal_rate is not None:
        lines.append(f"**Best refusal rate achieved:** {rec.best_refusal_rate:.1%}")
    if rec.best_coherence is not None:
        lines.append(f"**Best coherence achieved:** {rec.best_coherence:.3f}")

    lines.append("")
    lines.append(f"*{rec.reason}*")

    return "\n".join(lines)
