"""
QppgMonitor — Tier 1: Agent Monitoring SDK
==========================================

Drop-in wrapper around any LLM ReAct agent that scores every response
in real-time, fires alerts when risk exceeds a threshold, and accumulates
stats for dashboard/reporting.

Usage:
    from qppg_service import QppgMonitor

    monitor = QppgMonitor(threshold=0.65)

    # After each agent run:
    alert = monitor.track(question, steps, final_answer, finished=True)
    if alert:
        print(f"HIGH RISK: {alert.risk_score:.2f} — {alert.recommendation}")

    # Periodic stats:
    print(monitor.get_stats())
    monitor.export_report()

Zero cold-start: uses behavioral signals (SC2 step count) from the first query.
Calibration improves cross-domain accuracy after ~5 logged chains.
"""

from __future__ import annotations

import asyncio
import contextvars
import functools
import json
import time
import csv
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Optional

from .label_free_scorer import LabelFreeScorer, LabelFreeResult

if TYPE_CHECKING:
    from .alerting import WebhookAlerter


# ── Context variables for @monitor.watch decorator ───────────────────────────

#: Active step list for the current @monitor.watch frame; None = no active context.
_watch_steps: contextvars.ContextVar[list[dict] | None] = contextvars.ContextVar(
    "qppg_watch_steps", default=None
)


def add_step(
    thought: str = "",
    action_type: str = "Search",
    action_arg: str = "",
    observation: str = "",
) -> None:
    """Register one ReAct step inside a @monitor.watch-decorated function.

    Call this after every tool call so the monitor can score the chain when
    the decorated function returns.

    Example::

        @monitor.watch(domain="prod")
        async def run_agent(question: str) -> str:
            obs = await search_tool(question)
            add_step(thought="Searching...", action_type="Search",
                     action_arg=question, observation=obs)
            answer = await llm(obs)
            add_step(thought="Answering.", action_type="Finish",
                     action_arg=answer, observation="")
            return answer
    """
    steps = _watch_steps.get()
    if steps is not None:
        steps.append({
            "thought":     thought,
            "action_type": action_type,
            "action_arg":  action_arg,
            "observation": observation,
        })


@dataclass
class QppgAlert:
    """Fired when agent risk exceeds threshold."""
    timestamp:      float
    question:       str
    final_answer:   str
    risk_score:     float
    needs_review:   bool
    behavioral_score: float
    step_count:     int
    finished:       bool
    recommendation: str
    result:         LabelFreeResult = field(repr=False)


@dataclass
class MonitorStats:
    n_tracked:    int
    n_alerts:     int
    alert_rate:   float
    avg_risk:     float
    avg_steps:    float
    p95_risk:     float
    high_risk_q:  list[str]   # questions with highest risk


class QppgMonitor:
    """
    Real-time agent monitor.  Wraps a LabelFreeScorer and accumulates
    per-query history for reporting and calibration.

    Parameters
    ----------
    threshold       : risk score that triggers an alert (default 0.65)
    alert_callback  : optional callable(QppgAlert) called on every alert
    scorer          : pre-built LabelFreeScorer; if None a new one is created
    auto_calibrate  : if True, each logged chain is automatically fed into
                      scorer.calibrate() once n_logged >= min_cal_size
    min_cal_size    : minimum chains before auto-calibration triggers (default 5)
    domain          : optional string tag for multi-domain logging
    """

    def __init__(
        self,
        threshold: float = 0.65,
        alert_callback: Optional[Callable[[QppgAlert], None]] = None,
        scorer: Optional[LabelFreeScorer] = None,
        auto_calibrate: bool = True,
        min_cal_size: int = 5,
        domain: str = "default",
        *,
        db_path: Optional[str] = None,
        recal_every: int = 100,
        model_name: str = "",
        check_drift: bool = False,
        alerter: Optional["WebhookAlerter"] = None,
        online_decay: float = 0.2,
    ):
        self.threshold      = threshold
        self.alert_callback = alert_callback
        self.scorer         = scorer or LabelFreeScorer(review_threshold=threshold)
        self.auto_calibrate = auto_calibrate
        self.min_cal_size   = min_cal_size
        self.domain         = domain
        self.recal_every    = recal_every
        self.model_name     = model_name
        self.check_drift    = check_drift
        self._alerter       = alerter
        # Online GMM decay (exp69): 0<decay<1 blends new batch with old pool
        # decay=0.2 → best retention (0.694) vs catastrophic forgetting at 0.8
        # decay=1.0 → full retrain (legacy behavior)
        self.online_decay   = max(0.0, min(1.0, online_decay))

        self._history: list[dict] = []
        self._calibration_pool: list[dict] = []
        self._calibrated        = False
        self._n_cal_used        = 0
        self._n_since_last_drift_check = 0

        # SQLite persistence (optional)
        self._store = None
        if db_path is not None:
            from .store import ChainStore
            self._store = ChainStore(db_path)

        # Drift detector (optional)
        self._drift_detector = None
        if check_drift:
            from .drift import DriftDetector
            self._drift_detector = DriftDetector()

    # ── Public API ────────────────────────────────────────────────────────────

    def track(
        self,
        question: str,
        steps: list[dict],
        final_answer: str,
        finished: bool = True,
        metadata: Optional[dict] = None,
    ) -> Optional[QppgAlert]:
        """
        Score a completed agent run.  Returns a QppgAlert if risk >= threshold,
        else None.

        Parameters
        ----------
        question     : the original user question
        steps        : list of step dicts from the agent (same format as qppg_service)
        final_answer : the agent's final answer string
        finished     : True if agent called Finish[], False if it hit max steps
        metadata     : optional dict stored in history (run_id, user_id, etc.)
        """
        result = self.scorer.score(question, steps, final_answer, finished)
        step_count = sum(1 for s in steps if s.get("action_type") == "Search")

        entry = {
            "timestamp":       time.time(),
            "question":        question,
            "final_answer":    final_answer,
            "step_count":      step_count,
            "finished":        finished,
            "risk_score":      result.risk_score,
            "behavioral_score": result.behavioral_score,
            "needs_review":    result.needs_review,
            "domain":          self.domain,
            "metadata":        metadata or {},
            # store minimal chain for auto-calibration
            "_steps":          steps,
        }
        self._history.append(entry)
        self._calibration_pool.append({"question": question, "steps": steps})

        # Persist to SQLite if store configured
        if self._store is not None:
            self._store.add_chain(
                self.domain,
                {"question": question, "steps": steps,
                 "final_answer": final_answer, "finished": finished},
                risk_score=result.risk_score,
                alert=result.needs_review,
                failure_mode=result.failure_mode or "",
                model_name=self.model_name,
            )

        # Auto-calibrate once we have enough chains
        if (self.auto_calibrate and not self._calibrated
                and len(self._calibration_pool) >= self.min_cal_size):
            self._do_calibrate()

        # Progressive re-calibration every recal_every new chains
        if (self.auto_calibrate and self._calibrated
                and (len(self._calibration_pool) - self._n_cal_used) >= self.recal_every):
            self._do_calibrate()

        # Drift check every 50 tracked queries (if enabled and store available)
        self._n_since_last_drift_check += 1
        if (self.check_drift and self._drift_detector is not None
                and self._store is not None
                and self._n_since_last_drift_check >= 50):
            self._n_since_last_drift_check = 0
            drift = self._drift_detector.check(self.domain, self._store)
            if drift and self.alert_callback:
                # Re-use alert_callback with a sentinel alert for drift
                self.alert_callback(drift)  # type: ignore[arg-type]
            if drift and self._alerter is not None:
                self._alerter.fire_drift(drift)

        # Webhook: notify on alert-rate spike using last 100 chains
        if result.needs_review and self._alerter is not None:
            recent_n      = min(100, len(self._history))
            recent_alerts = sum(1 for e in self._history[-100:] if e.get("needs_review"))
            self._alerter.fire_alert_rate(self.domain, recent_alerts / recent_n, recent_n)

        if result.needs_review:
            alert = QppgAlert(
                timestamp        = entry["timestamp"],
                question         = question,
                final_answer     = final_answer,
                risk_score       = result.risk_score,
                needs_review     = True,
                behavioral_score = result.behavioral_score,
                step_count       = step_count,
                finished         = finished,
                recommendation   = self._make_recommendation(result, step_count, finished),
                result           = result,
            )
            if self.alert_callback:
                self.alert_callback(alert)
            return alert

        return None

    def track_batch(
        self,
        chains: list[dict],
        metadata_list: Optional[list[dict]] = None,
    ) -> list[Optional[QppgAlert]]:
        """
        Score a batch of chains at once.  Each chain dict must have keys:
        question, steps, final_answer, finished (optional).
        """
        alerts = []
        for i, chain in enumerate(chains):
            meta = metadata_list[i] if metadata_list else None
            alert = self.track(
                chain["question"],
                chain["steps"],
                chain.get("final_answer", ""),
                chain.get("finished", True),
                metadata=meta,
            )
            alerts.append(alert)
        return alerts

    def calibrate(self, chains: list[dict]) -> dict:
        """Manually provide calibration chains (no labels needed)."""
        result = self.scorer.calibrate(chains)
        self._calibrated = True
        self._n_cal_used = len(self._calibration_pool)
        return result

    def get_stats(self) -> MonitorStats:
        """Return aggregated monitoring statistics."""
        import numpy as np

        if not self._history:
            return MonitorStats(0, 0, 0.0, 0.0, 0.0, 0.0, [])

        risks  = [e["risk_score"]  for e in self._history]
        steps  = [e["step_count"]  for e in self._history]
        alerts = [e for e in self._history if e["needs_review"]]

        # Top 3 highest-risk questions
        sorted_h = sorted(self._history, key=lambda e: e["risk_score"], reverse=True)
        top3_q   = [e["question"][:80] + "…" if len(e["question"]) > 80
                    else e["question"]
                    for e in sorted_h[:3]]

        return MonitorStats(
            n_tracked   = len(self._history),
            n_alerts    = len(alerts),
            alert_rate  = len(alerts) / len(self._history),
            avg_risk    = float(np.mean(risks)),
            avg_steps   = float(np.mean(steps)),
            p95_risk    = float(np.percentile(risks, 95)),
            high_risk_q = top3_q,
        )

    def export_report(self) -> str:
        """Return a formatted text report of monitoring statistics."""
        s = self.get_stats()
        cal_status = (f"calibrated (n={self._n_cal_used} chains)"
                      if self._calibrated else
                      f"behavioral-only ({len(self._calibration_pool)}/{self.min_cal_size} chains for cal)")
        lines = [
            "=" * 60,
            "  QPPG Agent Monitor Report",
            "=" * 60,
            f"  Domain          : {self.domain}",
            f"  Scorer status   : {cal_status}",
            f"  Alert threshold : {self.threshold:.2f}",
            "",
            f"  Queries tracked : {s.n_tracked}",
            f"  Alerts fired    : {s.n_alerts}  ({s.alert_rate:.1%} of traffic)",
            f"  Avg risk score  : {s.avg_risk:.3f}",
            f"  P95 risk score  : {s.p95_risk:.3f}",
            f"  Avg search steps: {s.avg_steps:.1f}",
            "",
        ]
        if s.high_risk_q:
            lines.append("  Highest-risk queries:")
            for i, q in enumerate(s.high_risk_q, 1):
                lines.append(f"    {i}. {q}")
        lines.append("=" * 60)
        return "\n".join(lines)

    def export_csv(self, path: str | Path) -> None:
        """Export full history to CSV (for downstream analysis)."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        fields = ["timestamp", "question", "final_answer", "step_count",
                  "finished", "risk_score", "behavioral_score", "needs_review", "domain"]
        with path.open("w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=fields)
            w.writeheader()
            for e in self._history:
                w.writerow({k: e[k] for k in fields})

    def export_json(self, path: str | Path) -> None:
        """Export full history to JSON."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        exportable = [{k: v for k, v in e.items() if k != "_steps"}
                      for e in self._history]
        path.write_text(json.dumps(exportable, indent=2))

    def reset(self) -> None:
        """Clear history (keeps scorer calibration)."""
        self._history.clear()
        self._calibration_pool.clear()

    def watch(
        self,
        domain: str | None = None,
        finished: bool = True,
    ) -> Callable:
        """Decorator that automatically tracks an agent function's output.

        The decorated function must accept ``question: str`` as its first
        positional argument and return the final answer string.  Inside the
        function, call :func:`add_step` after each tool invocation so the
        monitor has the full reasoning chain.

        Parameters
        ----------
        domain   : override the monitor's default domain for this function.
        finished : treat the chain as a completed Finish[] call (default True).

        Example::

            from qppg_service import QppgMonitor, add_step

            monitor = QppgMonitor(threshold=0.65)

            @monitor.watch(domain="prod")
            async def run_agent(question: str) -> str:
                obs = await search(question)
                add_step(thought="Searching.", action_type="Search",
                         action_arg=question, observation=obs)
                answer = synthesise(obs)
                add_step(thought="Done.", action_type="Finish",
                         action_arg=answer, observation="")
                return answer

            # Alert is stored in run_agent.last_alert after each call.
        """
        def decorator(fn: Callable) -> Callable:
            if asyncio.iscoroutinefunction(fn):
                @functools.wraps(fn)
                async def async_wrapper(question: str, *args, **kwargs):
                    steps: list[dict] = []
                    token = _watch_steps.set(steps)
                    try:
                        final_answer = await fn(question, *args, **kwargs)
                    finally:
                        _watch_steps.reset(token)
                    _domain = domain or self.domain
                    old_domain, self.domain = self.domain, _domain
                    try:
                        alert = self.track(
                            question, steps, str(final_answer), finished
                        )
                    finally:
                        self.domain = old_domain
                    async_wrapper.last_alert = alert
                    return final_answer

                async_wrapper.last_alert = None
                return async_wrapper
            else:
                @functools.wraps(fn)
                def sync_wrapper(question: str, *args, **kwargs):
                    steps: list[dict] = []
                    token = _watch_steps.set(steps)
                    try:
                        final_answer = fn(question, *args, **kwargs)
                    finally:
                        _watch_steps.reset(token)
                    _domain = domain or self.domain
                    old_domain, self.domain = self.domain, _domain
                    try:
                        alert = self.track(
                            question, steps, str(final_answer), finished
                        )
                    finally:
                        self.domain = old_domain
                    sync_wrapper.last_alert = alert
                    return final_answer

                sync_wrapper.last_alert = None
                return sync_wrapper

        return decorator

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _do_calibrate(self) -> None:
        """
        (Re-)calibrate the scorer.

        First calibration: full fit on all available chains.
        Subsequent re-calibrations: online GMM blending (exp69) when
        online_decay < 1.0.  Decay-weighted sampling of old pool prevents
        catastrophic forgetting (0.694 retention at decay=0.2, exp69).
        """
        import numpy as np

        pool = self._calibration_pool
        # If SQLite store available, use persisted chains (capped at 1000)
        if self._store is not None:
            stored = self._store.get_calibration_pool(self.domain, n=1000)
            if stored:
                pool = stored

        if self._calibrated and 0.0 < self.online_decay < 1.0:
            pool = self._blend_pool_online(pool)

        self.scorer.calibrate(pool)
        self._calibrated = True
        self._n_cal_used = len(self._calibration_pool)
        if self._store is not None:
            auroc = self.scorer.status().get("expected_auroc", "unknown")
            try:
                auroc_f = float(str(auroc).split()[0].strip("~"))
            except (ValueError, IndexError):
                auroc_f = 0.0
            self._store.log_calibration(self.domain, len(pool), auroc_f)

    def _blend_pool_online(self, pool: list[dict]) -> list[dict]:
        """
        Return a decay-weighted blend of old and new chains (exp69, decay=0.2).

        New chains (last recal_every) get weight 1.0; older chains are
        down-weighted exponentially.  The blended set is capped at
        4 * recal_every chains to bound calibration cost.
        """
        import numpy as np

        n_new = min(self.recal_every, len(pool))
        new_batch = pool[-n_new:]
        old_pool  = pool[:-n_new]

        if not old_pool:
            return new_batch

        n_old = len(old_pool)
        # age=1 for most recent old chain, age=n_old for oldest
        ages = np.arange(n_old, 0, -1, dtype=float)
        weights = self.online_decay ** (ages / max(self.recal_every, 1))
        weights /= weights.sum()

        # Sample min(3*recal_every, n_old) old chains with decay weights
        n_sample = min(3 * self.recal_every, n_old)
        rng = np.random.default_rng(42)
        idx = rng.choice(n_old, size=n_sample, replace=False, p=weights)
        sampled_old = [old_pool[i] for i in sorted(idx)]

        return sampled_old + new_batch

    @staticmethod
    def _make_recommendation(result: LabelFreeResult, step_count: int, finished: bool) -> str:
        """Human-readable recommendation based on risk signals."""
        parts = []
        if not finished:
            parts.append("agent hit step limit — consider increasing max_steps")
        if step_count >= 5:
            parts.append(f"agent took {step_count} search steps (high) — likely struggling")
        if result.behavioral_score > 0.7:
            parts.append("high behavioral risk — review reasoning chain")
        if not parts:
            parts.append("general uncertainty — manual review recommended")
        return "; ".join(parts)


# ── Online RL: contextual bandit for signal weight adaptation ─────────────────

class OnlineBanditWeights:
    """
    Online RL (contextual bandit) for adapting signal weights from human feedback.

    Exp95 finding: Sonnet judge alone achieves ~0.97 optimal weight; SC_OLD and
    other signals shrink toward 0 when optimized offline. This bandit handles the
    online adaptation case where production feedback gradually shifts domain distribution.

    Usage (in production after deployment):
        bandit = OnlineBanditWeights(signal_names=["sc_old", "judge"])

        # After each alert, when a human reviews and confirms/denies:
        bandit.record_feedback(
            signal_scores=np.array([sc_score, judge_score]),
            was_true_positive=True,   # human confirmed this was a real failure
        )

        # To score a new chain using adapted weights:
        combined_risk = bandit.score(np.array([sc_score, judge_score]))
    """

    def __init__(
        self,
        signal_names: list[str] | None = None,
        n_signals: int = 2,
        lr: float = 0.02,
        seed: int = 42,
    ):
        import numpy as np
        self._np = np
        if signal_names is not None:
            n_signals = len(signal_names)
        self.signal_names = signal_names or [f"signal_{i}" for i in range(n_signals)]
        self.weights = np.ones(n_signals, dtype=float)
        self.lr = lr
        self.n_updates = 0
        self._reward_history: list[float] = []
        rng = np.random.default_rng(seed)
        self._rng = rng

    def score(self, signal_scores) -> float:
        """Return weighted combination of signal scores (inference)."""
        np = self._np
        signal_scores = np.asarray(signal_scores, dtype=float)
        w = np.abs(self.weights) / (np.abs(self.weights).sum() + 1e-9)
        return float(signal_scores @ w)

    def record_feedback(self, signal_scores, was_true_positive: bool) -> None:
        """
        REINFORCE-style update: reward=+1 for true positive (real failure caught),
        reward=-1 for false positive (correct chain flagged in error).

        Increase weights for signals that contributed to correct alerts;
        decrease weights that drove false positive alerts.
        """
        np = self._np
        signal_scores = np.asarray(signal_scores, dtype=float)
        reward = 1.0 if was_true_positive else -1.0
        w = np.abs(self.weights) / (np.abs(self.weights).sum() + 1e-9)
        score = max(float(signal_scores @ w), 1e-6)
        grad = reward * signal_scores / score
        self.weights += self.lr * grad
        self.weights = np.maximum(self.weights, 0.01)  # keep non-negative
        self.n_updates += 1
        self._reward_history.append(reward)

    def get_normalized_weights(self) -> dict[str, float]:
        """Return current weights normalized to sum=1."""
        np = self._np
        w = np.abs(self.weights) / np.abs(self.weights).sum()
        return {name: round(float(wi), 4) for name, wi in zip(self.signal_names, w)}

    def to_dict(self) -> dict:
        return {
            "weights": self.get_normalized_weights(),
            "n_updates": self.n_updates,
            "cumulative_reward": float(sum(self._reward_history)),
            "recent_precision": (
                float(sum(1 for r in self._reward_history[-20:] if r > 0) /
                      max(1, len(self._reward_history[-20:])))
            ),
        }
