"""
tracker.py — Metrics tracking, validation, and RL feedback loop for llm-guard-kit MCP server.

Three concerns:
  1. MetricsTracker  — SQLite store for every scored chain + human feedback
  2. AUROCTracker    — rolling AUROC over time windows (validation)
  3. RLFeedbackLoop  — online LogReg retraining from accumulated human labels
"""

from __future__ import annotations

import json
import math
import sqlite3
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np

DB_PATH = Path.home() / ".llm_guard_mcp" / "metrics.db"
DB_PATH.parent.mkdir(parents=True, exist_ok=True)


# ── Schema ─────────────────────────────────────────────────────────────────────

_SCHEMA = """
CREATE TABLE IF NOT EXISTS chains (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    ts          REAL    NOT NULL,           -- unix timestamp
    session_id  TEXT,
    question    TEXT    NOT NULL,
    final_answer TEXT,
    n_steps     INTEGER,
    risk_score  REAL,
    tier        TEXT,                       -- LOW / MEDIUM / HIGH
    ptrue_score REAL,
    beh_score   REAL,
    aborted     INTEGER DEFAULT 0,          -- 1 if stream_guard aborted
    human_label INTEGER DEFAULT NULL,       -- NULL=unlabeled, 0=correct, 1=wrong
    feedback_ts REAL    DEFAULT NULL,       -- when human labeled it
    source      TEXT    DEFAULT 'api',      -- api / mcp / cursor / claude_code
    metadata    TEXT    DEFAULT '{}'        -- JSON blob for extras
);

CREATE TABLE IF NOT EXISTS rl_checkpoints (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    ts          REAL    NOT NULL,
    n_labels    INTEGER,
    auroc       REAL,
    precision   REAL,
    recall      REAL,
    model_bytes BLOB                        -- pickled sklearn model
);

CREATE TABLE IF NOT EXISTS auroc_windows (
    ts          REAL    NOT NULL,
    window_days INTEGER NOT NULL,
    n_chains    INTEGER,
    n_labeled   INTEGER,
    auroc       REAL,
    PRIMARY KEY (ts, window_days)
);

CREATE TABLE IF NOT EXISTS heal_outcomes (
    id            INTEGER PRIMARY KEY AUTOINCREMENT,
    ts            REAL    NOT NULL,           -- unix timestamp
    chain_id      INTEGER NOT NULL,           -- FK into chains.id
    failure_mode  TEXT,                       -- e.g. RETRIEVAL_FAILURE
    action_type   TEXT,                       -- e.g. REPHRASE_QUERY
    success       INTEGER NOT NULL,           -- 1 = repair worked, 0 = still failed
    note          TEXT    DEFAULT ''          -- optional operator note
);
"""


# ── MetricsTracker ──────────────────────────────────────────────────────────────

@dataclass
class ChainRecord:
    question:     str
    final_answer: str
    n_steps:      int
    risk_score:   float
    tier:         str
    ptrue_score:  float  = 0.0
    beh_score:    float  = 0.0
    aborted:      bool   = False
    session_id:   str    = ""
    source:       str    = "api"
    metadata:     dict   = None

    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


class MetricsTracker:
    """Persistent SQLite store for all scored chains + human feedback."""

    def __init__(self, db_path: Path = DB_PATH):
        self.db_path = db_path
        self._con = sqlite3.connect(str(db_path), check_same_thread=False)
        self._con.row_factory = sqlite3.Row
        self._con.executescript(_SCHEMA)
        self._con.commit()

    # ── Write ─────────────────────────────────────────────────────────────────

    def log_chain(self, record: ChainRecord) -> int:
        """Insert a scored chain. Returns the new row ID."""
        cur = self._con.execute(
            """INSERT INTO chains
               (ts, session_id, question, final_answer, n_steps, risk_score, tier,
                ptrue_score, beh_score, aborted, source, metadata)
               VALUES (?,?,?,?,?,?,?,?,?,?,?,?)""",
            (
                time.time(),
                record.session_id,
                record.question[:500],
                (record.final_answer or "")[:500],
                record.n_steps,
                record.risk_score,
                record.tier,
                record.ptrue_score,
                record.beh_score,
                int(record.aborted),
                record.source,
                json.dumps(record.metadata),
            )
        )
        self._con.commit()
        return cur.lastrowid

    def add_feedback(self, chain_id: int, correct: bool) -> bool:
        """
        Record human feedback: correct=True means the answer WAS right.
        Returns True if the row was found and updated.
        """
        n = self._con.execute(
            "UPDATE chains SET human_label=?, feedback_ts=? WHERE id=?",
            (int(not correct), time.time(), chain_id)
        ).rowcount
        self._con.commit()
        return n > 0

    # ── Read ──────────────────────────────────────────────────────────────────

    def get_labeled(self, days: float = None) -> List[dict]:
        """Return all chains with human labels (optionally within last N days)."""
        q = "SELECT * FROM chains WHERE human_label IS NOT NULL"
        params = []
        if days is not None:
            q += " AND ts >= ?"
            params.append(time.time() - days * 86400)
        q += " ORDER BY ts DESC"
        return [dict(r) for r in self._con.execute(q, params)]

    def get_recent(self, n: int = 100) -> List[dict]:
        """Return last N scored chains (labeled or not)."""
        return [dict(r) for r in self._con.execute(
            "SELECT * FROM chains ORDER BY ts DESC LIMIT ?", (n,))]

    def summary(self, days: float = 7) -> dict:
        cutoff = time.time() - days * 86400
        rows = self._con.execute(
            "SELECT tier, risk_score, aborted, human_label FROM chains WHERE ts >= ?",
            (cutoff,)
        ).fetchall()
        total = len(rows)
        if total == 0:
            return {"total": 0, "days": days}
        labeled = [r for r in rows if r["human_label"] is not None]
        wrong   = [r for r in labeled if r["human_label"] == 1]
        return {
            "days":           days,
            "total":          total,
            "aborted":        sum(1 for r in rows if r["aborted"]),
            "tier_LOW":       sum(1 for r in rows if r["tier"] == "LOW"),
            "tier_MEDIUM":    sum(1 for r in rows if r["tier"] == "MEDIUM"),
            "tier_HIGH":      sum(1 for r in rows if r["tier"] == "HIGH"),
            "mean_risk":      round(sum(r["risk_score"] for r in rows) / total, 4),
            "n_labeled":      len(labeled),
            "n_wrong_labeled": len(wrong),
            "label_error_rate": round(len(wrong) / max(len(labeled), 1), 3),
        }

    def get_labeled_by_backend(self, backend: str, days: float = None) -> List[dict]:
        """Return labeled chains where metadata->judge_backend matches backend."""
        q = "SELECT * FROM chains WHERE human_label IS NOT NULL"
        params: list = []
        if days is not None:
            q += " AND ts >= ?"
            params.append(time.time() - days * 86400)
        q += " ORDER BY ts DESC"
        rows = [dict(r) for r in self._con.execute(q, params)]
        # Filter by judge_backend stored in metadata JSON
        result = []
        for r in rows:
            try:
                meta = json.loads(r.get("metadata") or "{}")
                if meta.get("judge_backend") == backend:
                    result.append(r)
            except Exception:
                pass
        return result

    def backend_summary(self, days: float = 30) -> dict:
        """Return per-backend chain counts and error rates."""
        cutoff = time.time() - days * 86400
        rows = self._con.execute(
            "SELECT metadata, human_label FROM chains WHERE ts >= ?", (cutoff,)
        ).fetchall()
        backends: dict = {}
        for r in rows:
            try:
                meta = json.loads(r["metadata"] or "{}")
                b = meta.get("judge_backend", "unknown")
            except Exception:
                b = "unknown"
            if b not in backends:
                backends[b] = {"total": 0, "labeled": 0, "wrong": 0}
            backends[b]["total"] += 1
            if r["human_label"] is not None:
                backends[b]["labeled"] += 1
                if r["human_label"] == 1:
                    backends[b]["wrong"] += 1
        for b, stats in backends.items():
            n = stats["labeled"]
            stats["error_rate"] = round(stats["wrong"] / n, 3) if n > 0 else None
        return {"days": days, "backends": backends}

    def domain_stats(self, days: float = 30) -> dict:
        """
        Per-domain failure rates from labeled chains.
        Domain is stored in metadata['domain']; falls back to 'general'.
        Returns {domain: {total, labeled, wrong, error_rate, nines}} sorted by error_rate desc.
        """
        import math
        cutoff = time.time() - days * 86400
        rows = self._con.execute(
            "SELECT metadata, human_label FROM chains WHERE ts >= ?", (cutoff,)
        ).fetchall()

        domains: dict = {}
        for r in rows:
            try:
                meta = json.loads(r["metadata"] or "{}")
                domain = meta.get("domain", "general")
            except Exception:
                domain = "general"
            if domain not in domains:
                domains[domain] = {"total": 0, "labeled": 0, "wrong": 0}
            domains[domain]["total"] += 1
            if r["human_label"] is not None:
                domains[domain]["labeled"] += 1
                if r["human_label"] == 1:
                    domains[domain]["wrong"] += 1

        # Compute error rates and nines for each domain
        for d, s in domains.items():
            n = s["labeled"]
            if n > 0:
                err = s["wrong"] / n
                s["error_rate"] = round(err, 3)
                success = 1.0 - err
                s["nines"] = round(-math.log10(max(1.0 - success, 1e-9)), 2) if success < 1.0 else None
            else:
                s["error_rate"] = None
                s["nines"] = None

        # Sort by error rate descending (worst domains first)
        sorted_domains = dict(
            sorted(domains.items(),
                   key=lambda x: x[1]["error_rate"] if x[1]["error_rate"] is not None else -1,
                   reverse=True)
        )
        return {"days": days, "domains": sorted_domains}

    def pending_feedback(self, limit: int = 10) -> List[dict]:
        """Return recently scored HIGH-risk chains awaiting human feedback."""
        return [dict(r) for r in self._con.execute(
            """SELECT id, ts, question, final_answer, risk_score, tier
               FROM chains WHERE human_label IS NULL AND tier='HIGH'
               ORDER BY ts DESC LIMIT ?""", (limit,))]

    def log_heal_outcome(
        self,
        chain_id:     int,
        success:      bool,
        failure_mode: str = "",
        action_type:  str = "",
        note:         str = "",
    ) -> int:
        """
        Record whether a SelfHealer repair attempt succeeded.

        Parameters
        ----------
        chain_id : int
            ID of the chain that was repaired (from log_chain).
        success : bool
            True if the repaired agent run answered correctly, False if still wrong.
        failure_mode : str
            FailureMode.primary_mode that triggered the heal (e.g. RETRIEVAL_FAILURE).
        action_type : str
            RecoveryAction.action_type used (e.g. REPHRASE_QUERY).
        note : str
            Optional operator annotation.

        Returns
        -------
        int : new heal_outcomes row ID.
        """
        cur = self._con.execute(
            """INSERT INTO heal_outcomes
               (ts, chain_id, failure_mode, action_type, success, note)
               VALUES (?,?,?,?,?,?)""",
            (time.time(), chain_id, failure_mode, action_type, int(success), note),
        )
        self._con.commit()
        return cur.lastrowid

    def heal_outcome_stats(self, days: float = 30) -> dict:
        """
        Aggregated SelfHealer effectiveness: success rate per failure_mode and action_type.

        Returns dict with overall success_rate and per-mode / per-action breakdown.
        """
        cutoff = time.time() - days * 86400
        rows = [dict(r) for r in self._con.execute(
            """SELECT failure_mode, action_type, success, COUNT(*) as n
               FROM heal_outcomes WHERE ts >= ?
               GROUP BY failure_mode, action_type, success""",
            (cutoff,),
        )]
        if not rows:
            return {"n_outcomes": 0, "success_rate": None, "by_mode": {}}

        total = sum(r["n"] for r in rows)
        success = sum(r["n"] for r in rows if r["success"])
        by_mode: dict = {}
        for r in rows:
            key = r["failure_mode"] or "unknown"
            if key not in by_mode:
                by_mode[key] = {"n": 0, "n_success": 0}
            by_mode[key]["n"] += r["n"]
            if r["success"]:
                by_mode[key]["n_success"] += r["n"]
        for k in by_mode:
            v = by_mode[k]
            v["success_rate"] = round(v["n_success"] / v["n"], 3) if v["n"] else None
        return {
            "n_outcomes":   total,
            "success_rate": round(success / total, 3) if total else None,
            "by_mode":      by_mode,
        }


# ── AUROCTracker ───────────────────────────────────────────────────────────────

class AUROCTracker:
    """Compute rolling AUROC over labeled feedback windows."""

    def __init__(self, tracker: MetricsTracker):
        self.tracker = tracker

    def compute(self, days: float = 30) -> Optional[float]:
        labeled = self.tracker.get_labeled(days=days)
        if len(labeled) < 10:
            return None
        try:
            from sklearn.metrics import roc_auc_score
            y     = np.array([r["human_label"] for r in labeled])
            score = np.array([r["risk_score"]  for r in labeled])
            if len(np.unique(y)) < 2:
                return None
            auroc = float(round(roc_auc_score(y, score), 4))
            # Store snapshot
            self.tracker._con.execute(
                "INSERT OR REPLACE INTO auroc_windows VALUES (?,?,?,?,?)",
                (time.time(), int(days), len(labeled), len(labeled), auroc)
            )
            self.tracker._con.commit()
            return auroc
        except Exception:
            return None

    def history(self, n: int = 20) -> List[dict]:
        """Return last N AUROC snapshots."""
        return [dict(r) for r in self.tracker._con.execute(
            "SELECT * FROM auroc_windows ORDER BY ts DESC LIMIT ?", (n,))]

    def compare_backends(self, days: float = 30) -> dict:
        """
        Compute AUROC per judge backend for A/B comparison (Anthropic vs NIM).
        Returns a dict of {backend: auroc} for backends with ≥ 10 labeled chains.
        """
        try:
            from sklearn.metrics import roc_auc_score
        except ImportError:
            return {"error": "sklearn required"}

        results = {}
        all_backends = set()

        labeled = self.tracker.get_labeled(days=days)
        for r in labeled:
            try:
                meta = json.loads(r.get("metadata") or "{}")
                all_backends.add(meta.get("judge_backend", "unknown"))
            except Exception:
                all_backends.add("unknown")

        for backend in all_backends:
            rows = self.tracker.get_labeled_by_backend(backend, days=days)
            if len(rows) < 10:
                results[backend] = {"auroc": None, "n": len(rows), "note": "need ≥10 labels"}
                continue
            y = np.array([r["human_label"] for r in rows])
            scores = np.array([r["risk_score"] for r in rows])
            if len(np.unique(y)) < 2:
                results[backend] = {"auroc": None, "n": len(rows), "note": "single class"}
                continue
            auroc = float(round(roc_auc_score(y, scores), 4))
            results[backend] = {"auroc": auroc, "n": len(rows)}

        return {"days": days, "backends": results}

    def nines_dashboard(self, days: float = 30) -> dict:
        """
        Operational reliability dashboard: converts labeled success rates to
        Karpathy's 'march of nines' (nines = -log10(failure_rate)).

        Shows: current nines achieved, success rate, per-tier breakdown,
        rolling trend, and how many more labeled chains to the next nine.
        """
        import math

        labeled = self.tracker.get_labeled(days=days)
        n = len(labeled)

        if n < 5:
            return {
                "nines": None,
                "success_rate": None,
                "n_labeled": n,
                "status": f"Need at least 5 labeled chains (have {n}). Submit feedback via submit_feedback.",
                "interpretation": "Cannot compute nines yet.",
            }

        n_correct = sum(1 for r in labeled if r["human_label"] == 0)
        n_wrong   = sum(1 for r in labeled if r["human_label"] == 1)
        success_rate = n_correct / n

        # Nines: 90%→1, 99%→2, 99.9%→3, 99.99%→4
        failure_rate = 1.0 - success_rate
        nines = round(-math.log10(max(failure_rate, 1e-9)), 2) if failure_rate > 0 else None

        # Per-tier breakdown (how well does each tier correlate with actual correctness)
        tier_stats: dict = {}
        for r in labeled:
            t = r["tier"] or "UNKNOWN"
            if t not in tier_stats:
                tier_stats[t] = {"total": 0, "correct": 0}
            tier_stats[t]["total"] += 1
            if r["human_label"] == 0:
                tier_stats[t]["correct"] += 1
        tier_breakdown = {}
        for t, s in tier_stats.items():
            sr = s["correct"] / s["total"]
            tier_breakdown[t] = {
                "n": s["total"],
                "success_rate": round(sr, 3),
                "nines": round(-math.log10(max(1.0 - sr, 1e-9)), 2) if sr < 1.0 else None,
            }

        # How many chains needed to reach next nine?
        next_nine_target = math.ceil(nines + 1) if nines is not None else 1
        next_nine_success = 1.0 - (10 ** -next_nine_target)
        chains_needed = None
        if success_rate < next_nine_success and success_rate > 0:
            # Solve: (n_correct + x) / (n + x) = next_nine_success
            # x = (n * next_nine_success - n_correct) / (1 - next_nine_success)
            denom = 1.0 - next_nine_success
            if denom > 0:
                x = (n * next_nine_success - n_correct) / denom
                chains_needed = max(0, int(math.ceil(x)))

        # Rolling trend: nines over last 3 AUROC snapshots
        hist = self.history(5)
        nines_trend = []
        for h in hist:
            if h["auroc"] is not None:
                # Approximate nines from AUROC — rough proxy
                # AUROC 0.75 ≈ catches 75% of failures → effective nines improvement
                # Use actual labeled chains per window where possible
                nines_trend.append({"ts": h["ts"], "auroc": h["auroc"]})

        # Karpathy interpretation
        if nines is None:
            interp = "No data yet."
        elif nines < 1.0:
            interp = f"{nines:.2f} nines ({success_rate*100:.1f}% success). Below 1 nine — demo quality. Most agents fail regularly."
        elif nines < 2.0:
            interp = f"{nines:.2f} nines ({success_rate*100:.1f}% success). 1 nine achieved. Viable for low-stakes use. Still 1 in {int(round(1/failure_rate))} chains wrong."
        elif nines < 3.0:
            interp = f"{nines:.2f} nines ({success_rate*100:.1f}% success). 2 nines. Production-viable with guard rails. 1 in {int(round(1/failure_rate))} wrong."
        else:
            interp = f"{nines:.2f} nines ({success_rate*100:.1f}% success). Excellent. 1 in {int(round(1/failure_rate))} wrong."

        return {
            "nines":               nines,
            "success_rate":        round(success_rate, 4),
            "n_labeled":           n,
            "n_correct":           n_correct,
            "n_wrong":             n_wrong,
            "days":                days,
            "status":              interp,
            "tier_breakdown":      tier_breakdown,
            "next_nine_target":    next_nine_target,
            "chains_to_next_nine": chains_needed,
            "auroc_trend":         nines_trend,
            "reference": {
                "1_nine":  "90.0% success — demo quality",
                "2_nines": "99.0% success — production viable",
                "3_nines": "99.9% success — mission critical",
                "4_nines": "99.99% success — infrastructure grade",
            },
        }

    def drift_alert(self, days: float = 30, min_delta: float = 0.05) -> Optional[str]:
        """
        Return a warning string if AUROC has dropped by min_delta from peak.
        Returns None if no drift detected or insufficient data.
        """
        hist = self.history(10)
        if len(hist) < 3:
            return None
        aurocs = [h["auroc"] for h in hist if h["auroc"] is not None]
        if len(aurocs) < 3:
            return None
        peak    = max(aurocs)
        current = aurocs[0]
        if peak - current >= min_delta:
            return (f"DRIFT ALERT: AUROC dropped from {peak:.3f} to {current:.3f} "
                    f"(Δ={peak-current:.3f} ≥ {min_delta} threshold). "
                    f"Consider retraining or reviewing recent flagged chains.")
        return None


# ── RLFeedbackLoop ─────────────────────────────────────────────────────────────

class RLFeedbackLoop:
    """
    Online supervised learning loop: retrain from accumulated human feedback labels.

    This is NOT classical RL (no policy network, no reward shaping).
    It is supervised retraining triggered by label accumulation, termed "RL" because
    human feedback acts as a reward signal that updates the model weights.

    Three learning phases keyed on label count:

    Phase 2 (≥30 labels): LogReg on 5 aggregate scores (risk, beh, ptrue, n_steps, tier).
        Improves calibration. Does NOT improve raw discrimination — same feature
        ceiling as SC_OLD. AUROC gain: 0-3pp on strong-signal domains.

    Phase 3 (≥100 labels per domain): Per-domain LogReg on all available signals.
        For each domain with ≥100 labels, selects best signal via 5-fold CV.
        Matches or slightly exceeds the static DomainRouter routing for that domain.
        Use per_domain_retrain() to access this phase.

    Phase 4 (≥500 labels): MiniJudge fine-tuning — NOT YET IMPLEMENTED.
        Cross-architecture disagreement trainable threshold — NOT YET IMPLEMENTED.
        Text-level chain reading (2000+ labels) — NOT YET IMPLEMENTED.

    Label collection:
        Currently: explicit human feedback via submit_feedback MCP tool only.
        NOT YET IMPLEMENTED: passive/implicit labels, A/B testing, cross-customer.

    Trigger: retrain when N new labels have accumulated since last checkpoint.
    """

    RETRAIN_EVERY      = 20    # retrain after every 20 new human labels
    MIN_LABELS         = 30    # minimum labels before first retrain
    MIN_PER_CLASS      = 5     # need at least 5 wrong and 5 correct
    PHASE3_MIN_DOMAIN  = 100   # minimum per-domain labels for Phase 3

    def __init__(self, tracker: MetricsTracker):
        self.tracker  = tracker
        self._clf     = None   # current global LogReg model (Phase 2)
        self._domain_clfs: dict = {}   # {domain: clf} for Phase 3 per-domain models
        self._n_at_last_train = 0

    def should_retrain(self) -> Tuple[bool, str]:
        labeled = self.tracker.get_labeled()
        n = len(labeled)
        n_wrong   = sum(1 for r in labeled if r["human_label"] == 1)
        n_correct = n - n_wrong
        if n < self.MIN_LABELS:
            return False, f"Need {self.MIN_LABELS - n} more labels before first retrain (have {n})"
        if n_wrong < self.MIN_PER_CLASS or n_correct < self.MIN_PER_CLASS:
            return False, f"Need ≥{self.MIN_PER_CLASS} per class (have {n_correct} correct, {n_wrong} wrong)"
        new_since_last = n - self._n_at_last_train
        if new_since_last < self.RETRAIN_EVERY:
            return False, f"Need {self.RETRAIN_EVERY - new_since_last} more labels before retraining"
        return True, f"Ready: {n} labels ({n_correct} correct, {n_wrong} wrong), {new_since_last} new since last train"

    def retrain(self, guard=None) -> dict:
        """
        Retrain LogReg on all labeled chains.
        Returns a dict with new AUROC, precision, recall, n_labels.
        If guard is provided, update guard's local_verifier weights.
        """
        from sklearn.linear_model import LogisticRegression
        from sklearn.preprocessing import StandardScaler
        from sklearn.pipeline import Pipeline
        from sklearn.model_selection import cross_val_score
        from sklearn.metrics import roc_auc_score, precision_score, recall_score
        import pickle

        labeled = self.tracker.get_labeled()
        if not labeled:
            return {"error": "No labeled data"}

        # Features: [risk_score, beh_score, ptrue_score, n_steps, tier_num]
        def extract(r):
            tier_num = {"LOW": 0, "MEDIUM": 0.5, "HIGH": 1.0}.get(r["tier"], 0.5)
            return [
                r["risk_score"]  or 0.0,
                r["beh_score"]   or 0.0,
                r["ptrue_score"] or 0.0,
                min((r["n_steps"] or 3) / 10.0, 1.0),
                tier_num,
            ]

        X = np.array([extract(r) for r in labeled])
        y = np.array([r["human_label"] for r in labeled])

        if len(np.unique(y)) < 2:
            return {"error": "Single class — need more labels of the other class"}

        clf = Pipeline([
            ("sc", StandardScaler()),
            ("lr", LogisticRegression(max_iter=500, class_weight="balanced",
                                       C=1.0, solver="lbfgs"))
        ])

        # Cross-validated AUROC (honest estimate)
        if len(y) >= 20:
            cv_aurocs = cross_val_score(clf, X, y, cv=min(5, len(y)//5),
                                         scoring="roc_auc")
            cv_auroc = float(round(np.mean(cv_aurocs), 4))
        else:
            cv_auroc = None

        # Fit on all data
        clf.fit(X, y)
        self._clf = clf

        proba = clf.predict_proba(X)[:, 1]
        train_auroc = float(round(roc_auc_score(y, proba), 4))

        # Precision/recall at 0.5 threshold
        pred = (proba >= 0.5).astype(int)
        prec = float(round(precision_score(y, pred, zero_division=0), 3))
        rec  = float(round(recall_score(y, pred,    zero_division=0), 3))

        self._n_at_last_train = len(labeled)

        # Persist checkpoint
        model_bytes = pickle.dumps(clf)
        self.tracker._con.execute(
            "INSERT INTO rl_checkpoints (ts, n_labels, auroc, precision, recall, model_bytes) VALUES (?,?,?,?,?,?)",
            (time.time(), len(labeled), train_auroc, prec, rec, model_bytes)
        )
        self.tracker._con.commit()

        result = {
            "status":       "retrained",
            "n_labels":     len(labeled),
            "n_wrong":      int(y.sum()),
            "n_correct":    int((1 - y).sum()),
            "train_auroc":  train_auroc,
            "cv_auroc":     cv_auroc,
            "precision":    prec,
            "recall":       rec,
            "improvement_note": (
                f"CV AUROC {cv_auroc:.3f}" if cv_auroc else
                f"Train AUROC {train_auroc:.3f} (need ≥20 labels for CV)"
            ),
        }

        # NOTE: We do NOT inject into guard._local_verifier here.
        # LocalVerifier expects 12 raw SC features (question, steps, final_answer).
        # RLFeedbackLoop uses 5 aggregated scores (risk, beh, ptrue, n_steps, tier).
        # These are different feature spaces — merging them would silently corrupt scoring.
        # Instead, rl_score is returned as a standalone signal in score_chain responses
        # and used to gate tier/needs_alert when available (see server.py).
        if guard is not None:
            result["guard_updated"] = False
            result["guard_note"] = (
                "RL model improves risk scoring via rl_score field. "
                "LocalVerifier (12 SC features) is separate and unchanged."
            )

        return result

    def predict(self, risk_score: float, beh_score: float, ptrue_score: float,
                n_steps: int, tier: str) -> Optional[float]:
        """Use current RL model to re-score a chain. Returns None if no model."""
        if self._clf is None:
            return None
        tier_num = {"LOW": 0, "MEDIUM": 0.5, "HIGH": 1.0}.get(tier, 0.5)
        X = np.array([[risk_score, beh_score, ptrue_score,
                       min(n_steps / 10.0, 1.0), tier_num]])
        return float(round(self._clf.predict_proba(X)[0, 1], 4))

    def load_latest(self) -> bool:
        """Load the most recent RL checkpoint from DB."""
        import pickle
        row = self.tracker._con.execute(
            "SELECT model_bytes, n_labels FROM rl_checkpoints ORDER BY ts DESC LIMIT 1"
        ).fetchone()
        if row and row["model_bytes"]:
            try:
                self._clf = pickle.loads(row["model_bytes"])
                self._n_at_last_train = row["n_labels"]
                return True
            except Exception:
                pass
        return False

    def checkpoint_history(self, n: int = 10) -> List[dict]:
        return [dict(r) for r in self.tracker._con.execute(
            "SELECT ts, n_labels, auroc, precision, recall FROM rl_checkpoints ORDER BY ts DESC LIMIT ?",
            (n,))]

    # ── Phase 3: per-domain signal selection (≥100 labels per domain) ──────────

    def per_domain_retrain(self) -> dict:
        """
        Phase 3 learning: fit a per-domain LogReg when ≥100 labels exist per domain.

        Uses the same 5 aggregate features as the global model, but trains separately
        per domain so the weights can specialize (e.g., HP domain may weight n_steps
        more heavily than TV domain).

        This does NOT improve the feature ceiling — it improves calibration and
        signal weighting per domain. Expect 0-5pp AUROC gain over the global model
        on domains with >100 labels.

        Returns a report of which domains were trained, their CV AUROCs, and which
        domains still lack enough data.
        """
        from sklearn.linear_model import LogisticRegression
        from sklearn.preprocessing import StandardScaler
        from sklearn.pipeline import Pipeline
        from sklearn.model_selection import cross_val_score
        import json

        labeled = self.tracker.get_labeled()
        if not labeled:
            return {"error": "No labeled data"}

        def extract(r):
            tier_num = {"LOW": 0, "MEDIUM": 0.5, "HIGH": 1.0}.get(r["tier"], 0.5)
            return [
                r["risk_score"]  or 0.0,
                r["beh_score"]   or 0.0,
                r["ptrue_score"] or 0.0,
                min((r["n_steps"] or 3) / 10.0, 1.0),
                tier_num,
            ]

        # Group by domain (stored in metadata JSON)
        domain_data: dict = {}
        for r in labeled:
            try:
                meta = json.loads(r.get("metadata") or "{}")
                domain = meta.get("domain", "general")
            except Exception:
                domain = "general"
            if domain not in domain_data:
                domain_data[domain] = {"X": [], "y": []}
            domain_data[domain]["X"].append(extract(r))
            domain_data[domain]["y"].append(r["human_label"])

        results = {"trained": [], "skipped": [], "phase": 3}

        for domain, data in domain_data.items():
            X = np.array(data["X"])
            y = np.array(data["y"])
            n = len(y)
            n_wrong = int(y.sum())
            n_correct = n - n_wrong

            if n < self.PHASE3_MIN_DOMAIN:
                results["skipped"].append({
                    "domain": domain,
                    "n": n,
                    "reason": f"Need {self.PHASE3_MIN_DOMAIN} labels (have {n})",
                })
                continue

            if n_wrong < self.MIN_PER_CLASS or n_correct < self.MIN_PER_CLASS:
                results["skipped"].append({
                    "domain": domain,
                    "n": n,
                    "reason": f"Need ≥{self.MIN_PER_CLASS} per class (have {n_correct} correct, {n_wrong} wrong)",
                })
                continue

            clf = Pipeline([
                ("sc", StandardScaler()),
                ("lr", LogisticRegression(max_iter=500, class_weight="balanced",
                                           C=1.0, solver="lbfgs"))
            ])
            cv_aurocs = cross_val_score(clf, X, y, cv=min(5, n // 20),
                                        scoring="roc_auc")
            cv_auroc = float(round(np.mean(cv_aurocs), 4))
            clf.fit(X, y)
            self._domain_clfs[domain] = clf

            results["trained"].append({
                "domain":    domain,
                "n":         n,
                "n_wrong":   n_wrong,
                "cv_auroc":  cv_auroc,
            })

        results["n_domain_models"] = len(self._domain_clfs)
        results["note"] = (
            "Phase 3 trains per-domain models. Feature ceiling unchanged from Phase 2. "
            "Expected gain: 0-5pp AUROC vs global model on calibration."
        )
        return results

    def predict_domain(self, risk_score: float, beh_score: float, ptrue_score: float,
                       n_steps: int, tier: str, domain: str = "general") -> Optional[float]:
        """
        Re-score using the domain-specific Phase 3 model if available,
        falling back to the global Phase 2 model, then None.
        """
        tier_num = {"LOW": 0, "MEDIUM": 0.5, "HIGH": 1.0}.get(tier, 0.5)
        X = np.array([[risk_score, beh_score, ptrue_score,
                       min(n_steps / 10.0, 1.0), tier_num]])
        clf = self._domain_clfs.get(domain) or self._clf
        if clf is None:
            return None
        return float(round(clf.predict_proba(X)[0, 1], 4))

    def phase_summary(self) -> dict:
        """
        Return a human-readable summary of current learning phase and what's needed
        to advance to the next phase.

        Phases:
          Phase 1 (0 labels):   DomainRouter pre-validated signals only.
          Phase 2 (≥30 labels): Global LogReg calibration. Trains automatically.
          Phase 3 (≥100/domain): Per-domain LogReg. Call per_domain_retrain().
          Phase 4 (≥500 labels): MiniJudge fine-tuning. NOT YET IMPLEMENTED.
        """
        labeled = self.tracker.get_labeled()
        n_total = len(labeled)
        n_wrong = sum(1 for r in labeled if r["human_label"] == 1)

        import json
        domain_counts: dict = {}
        for r in labeled:
            try:
                meta = json.loads(r.get("metadata") or "{}")
                d = meta.get("domain", "general")
            except Exception:
                d = "general"
            domain_counts[d] = domain_counts.get(d, 0) + 1

        domains_at_phase3 = [d for d, c in domain_counts.items() if c >= self.PHASE3_MIN_DOMAIN]

        if n_total < self.MIN_LABELS:
            current_phase = 1
            status = f"Phase 1 (pre-validated signals). Need {self.MIN_LABELS - n_total} more labels for Phase 2."
        elif n_total < self.PHASE3_MIN_DOMAIN:
            current_phase = 2
            status = f"Phase 2 (global calibration, n={n_total}). Need {self.PHASE3_MIN_DOMAIN - n_total} more for Phase 3."
        elif domains_at_phase3:
            current_phase = 3
            status = f"Phase 3 available for domains: {domains_at_phase3}. Call per_domain_retrain()."
        else:
            current_phase = 2
            status = f"Phase 2 (n={n_total}). No domain has ≥{self.PHASE3_MIN_DOMAIN} labels yet."

        return {
            "current_phase":    current_phase,
            "n_total_labels":   n_total,
            "n_wrong":          n_wrong,
            "n_correct":        n_total - n_wrong,
            "domain_counts":    domain_counts,
            "domains_at_phase3": domains_at_phase3,
            "phase2_active":    self._clf is not None,
            "phase3_domains":   list(self._domain_clfs.keys()),
            "status":           status,
            "phase4_note":      "Phase 4 (MiniJudge fine-tuning, ≥500 labels) — NOT YET IMPLEMENTED",
            "phase5_note":      "Phase 5 (text-level learning, ≥2000 labels) — NOT YET IMPLEMENTED",
            "passive_labels_note": "Passive/implicit label collection — NOT YET IMPLEMENTED",
        }
