"""
server.py — llm-guard-kit MCP Server.

Exposes llm-guard-kit as MCP tools that Claude, Cursor, and any MCP-compatible
host can call. Includes metrics tracking and RL feedback loop.

Tools exposed:
  score_chain         — score a completed agent chain (behavioral + P(True))
  stream_check        — mid-chain abort check (call after step 2)
  submit_feedback     — human labels a chain as correct/wrong (RL signal)
  get_metrics         — summary stats for a time window
  get_auroc           — rolling AUROC over labeled feedback
  trigger_retrain     — retrain RL model from accumulated labels
  get_pending_review  — chains awaiting human feedback
  get_rl_status       — current RL model status and training history
  compare_backends    — per-backend AUROC (Anthropic vs NIM)
  get_nines_dashboard — operational reliability in Karpathy's 'march of nines'
  get_domain_stats    — per-domain failure rates (which question types fail most)
  check_retry_budget  — retry vs stop decision for multi-attempt agent loops
  retrieval_audit     — per-step cosine similarity audit for RETRIEVAL_FAILURE diagnosis
  heal_outcome        — record whether a SelfHealer repair attempt succeeded (RL loop closure)
  get_heal_stats      — aggregated SelfHealer effectiveness per failure_mode

Run:
  python -m llm_guard_mcp.server          # stdio transport (Claude Desktop / Cursor)
  python -m llm_guard_mcp.server --sse    # SSE transport (web clients)
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

# ── Path setup ────────────────────────────────────────────────────────────────
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

from llm_guard.scoring.agent_guard import AgentGuard
from llm_guard_mcp.tracker import MetricsTracker, AUROCTracker, RLFeedbackLoop, ChainRecord

# ── Globals ───────────────────────────────────────────────────────────────────
_API_KEY     = os.environ.get("ANTHROPIC_API_KEY", "")
_NIM_API_KEY = os.environ.get("NIM_API_KEY", "")
_NIM_BASE_URL = os.environ.get("NIM_BASE_URL", "https://integrate.api.nvidia.com/v1")
_NIM_JUDGE_MODEL = os.environ.get("NIM_JUDGE_MODEL", "meta/llama-3.3-70b-instruct")
_NIM_PTRUE_MODEL = os.environ.get("NIM_PTRUE_MODEL", "meta/llama-3.1-8b-instruct")
_guard   = None
_tracker = MetricsTracker()
_auroc   = AUROCTracker(_tracker)
_rl      = RLFeedbackLoop(_tracker)
_rl.load_latest()   # restore last RL checkpoint if available


def _get_guard() -> AgentGuard:
    global _guard
    if _guard is None:
        has_key = bool(_API_KEY) or bool(_NIM_API_KEY)
        _guard = AgentGuard(
            api_key=_API_KEY if _API_KEY else None,
            use_judge=has_key,
            nim_api_key=_NIM_API_KEY if _NIM_API_KEY else None,
            nim_base_url=_NIM_BASE_URL,
            nim_judge_model=_NIM_JUDGE_MODEL,
            nim_ptrue_model=_NIM_PTRUE_MODEL,
        )
    return _guard


# ── MCP Server ────────────────────────────────────────────────────────────────
app = Server("llm-guard-kit")


@app.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="score_chain",
            description=(
                "Score a completed agent chain for reliability. "
                "Returns risk_score (0=safe, 1=wrong), tier (LOW/MEDIUM/HIGH), "
                "and whether human review is needed. "
                "Call this after every multi-step agent response before returning to user."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "question":     {"type": "string", "description": "The original user question"},
                    "steps":        {"type": "array",  "description": "List of agent steps. Each step: {thought, action_type, action_arg, observation}"},
                    "final_answer": {"type": "string", "description": "The agent's final answer"},
                    "session_id":   {"type": "string", "description": "Optional session identifier for tracking"},
                    "source":       {"type": "string", "description": "Source: api | cursor | claude_code | langchain", "default": "mcp"},
                    "domain":       {"type": "string", "description": "Optional query domain for failure tracking (e.g. 'medical', 'legal', 'finance'). Auto-detected if omitted."},
                },
                "required": ["question", "steps", "final_answer"],
            }
        ),
        Tool(
            name="stream_check",
            description=(
                "Mid-chain abort check. Call after step 2 (index 1) with steps so far. "
                "If should_abort=true, stop the agent immediately — it is unlikely to succeed. "
                "Saves API cost by aborting failing chains early."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "question":    {"type": "string"},
                    "steps_so_far": {"type": "array", "description": "Steps completed so far (typically 2)"},
                    "session_id":  {"type": "string"},
                },
                "required": ["question", "steps_so_far"],
            }
        ),
        Tool(
            name="submit_feedback",
            description=(
                "Submit human feedback on a scored chain. "
                "This is the RL signal: every label improves the model over time. "
                "Call this when a user marks an answer as correct or wrong."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "chain_id": {"type": "integer", "description": "ID returned by score_chain"},
                    "correct":  {"type": "boolean", "description": "true if the answer was correct, false if wrong"},
                    "note":     {"type": "string",  "description": "Optional note explaining the error"},
                },
                "required": ["chain_id", "correct"],
            }
        ),
        Tool(
            name="get_metrics",
            description="Get summary metrics for recent chains (default: last 7 days).",
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {"type": "number", "description": "Time window in days", "default": 7},
                },
            }
        ),
        Tool(
            name="get_auroc",
            description=(
                "Compute rolling AUROC from human-labeled feedback. "
                "Returns current model quality and drift alerts. "
                "Requires at least 10 labeled chains."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {"type": "number", "description": "Window for AUROC computation", "default": 30},
                },
            }
        ),
        Tool(
            name="trigger_retrain",
            description=(
                "Retrain the RL model from accumulated human feedback labels. "
                "Automatically triggered when 20+ new labels accumulate, "
                "but can be called manually. Returns new AUROC and precision/recall."
            ),
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="get_pending_review",
            description=(
                "Get HIGH-risk chains that are waiting for human feedback. "
                "Use this to show reviewers what needs labeling to improve the RL model."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "limit": {"type": "integer", "description": "Max chains to return", "default": 10},
                },
            }
        ),
        Tool(
            name="get_rl_status",
            description="Get current RL model status: training history, AUROC trend, drift alerts.",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="compare_backends",
            description=(
                "Compare AUROC by judge backend (Anthropic vs NVIDIA NIM vs behavioral). "
                "Shows per-backend accuracy from human-labeled chains. "
                "Use this to validate whether NIM improves scoring quality over time."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {"type": "number", "description": "Lookback window in days", "default": 30},
                },
            }
        ),
        Tool(
            name="get_nines_dashboard",
            description=(
                "Karpathy 'march of nines' reliability dashboard. "
                "Converts labeled success rates into nines metric (1 nine=90%, 2=99%, 3=99.9%). "
                "Shows current reliability tier, chains-to-next-nine projection, and per-tier breakdown. "
                "Call this to understand how many more labels/improvements are needed to reach the next reliability milestone."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {"type": "number", "description": "Lookback window in days", "default": 30},
                },
            }
        ),
        Tool(
            name="get_domain_stats",
            description=(
                "Per-domain failure rates from labeled chains. "
                "Shows which question types (person_factual, temporal, numeric, reasoning, etc.) fail most often. "
                "Use this to find where the agent is weakest and prioritize improvements."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {"type": "number", "description": "Lookback window in days", "default": 30},
                },
            }
        ),
        Tool(
            name="check_retry_budget",
            description=(
                "Retry vs stop decision for multi-attempt agent loops. "
                "Pass the risk scores from each attempt so far; returns action: "
                "'use_result' (good enough), 'retry' (improving, keep going), or "
                "'accept_uncertainty' (budget exhausted or stalling). "
                "Use this in agentic loops to avoid burning API budget on hopeless chains."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "attempt_risks": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "Risk scores from each attempt so far (e.g. [0.72, 0.65, 0.58])"
                    },
                    "max_retries": {"type": "integer", "description": "Max retries allowed", "default": 3},
                    "low_threshold": {"type": "number", "description": "Risk below this → use_result immediately", "default": 0.40},
                    "high_threshold": {"type": "number", "description": "Risk above this → still HIGH risk", "default": 0.65},
                },
                "required": ["attempt_risks"],
            }
        ),
        Tool(
            name="heal_outcome",
            description=(
                "Record whether a SelfHealer repair attempt succeeded. "
                "This closes the detect → diagnose → repair → outcome loop. "
                "Call this after running the repaired agent chain to record if it answered correctly. "
                "Outcomes are aggregated into success-rate stats per failure mode, "
                "enabling data-driven repair strategy selection over time."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "chain_id":     {"type": "integer", "description": "ID of the original chain (from score_chain response)"},
                    "success":      {"type": "boolean", "description": "true if the repaired run answered correctly, false if still wrong"},
                    "failure_mode": {"type": "string",  "description": "FailureMode that triggered the heal (e.g. RETRIEVAL_FAILURE)"},
                    "action_type":  {"type": "string",  "description": "SelfHealer action used (e.g. REPHRASE_QUERY, FORCE_FINISH)"},
                    "note":         {"type": "string",  "description": "Optional operator annotation"},
                },
                "required": ["chain_id", "success"],
            }
        ),
        Tool(
            name="get_heal_stats",
            description=(
                "Aggregated SelfHealer repair effectiveness: success rate per failure_mode. "
                "Shows which repair strategies work best and which failure modes are hardest to fix. "
                "Requires heal_outcome calls to accumulate data."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {"type": "number", "description": "Lookback window in days", "default": 30},
                },
            }
        ),
        Tool(
            name="retrieval_audit",
            description=(
                "Audit the retrieval quality of a completed agent chain. "
                "Returns per-step cosine similarity between the question and each "
                "search observation, plus quality_label (GOOD/FAIR/POOR) and aggregate stats. "
                "Correct agents average mean_sim=0.568, incorrect 0.441 (Δ=0.127, exp43). "
                "Use this to diagnose RETRIEVAL_FAILURE mode and identify which search steps "
                "returned off-topic results. Requires no API key — $0 per call."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "question": {
                        "type": "string",
                        "description": "The original question asked of the agent"
                    },
                    "steps": {
                        "type": "array",
                        "description": "List of agent steps (each with thought, action_type, action_arg, observation)",
                        "items": {
                            "type": "object",
                            "properties": {
                                "thought":      {"type": "string"},
                                "action_type":  {"type": "string"},
                                "action_arg":   {"type": "string"},
                                "observation":  {"type": "string"},
                            }
                        }
                    },
                },
                "required": ["question", "steps"],
            }
        ),
    ]


@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    try:
        result = await _dispatch(name, arguments)
        return [TextContent(type="text", text=json.dumps(result, indent=2))]
    except Exception as e:
        return [TextContent(type="text", text=json.dumps({
            "error": str(e), "tool": name
        }))]


async def _dispatch(name: str, args: dict) -> Any:

    # ── score_chain ────────────────────────────────────────────────────────────
    if name == "score_chain":
        guard        = _get_guard()
        question     = args["question"]
        steps        = args.get("steps", [])
        final_answer = args.get("final_answer", "")
        session_id   = args.get("session_id", "")
        source       = args.get("source", "mcp")

        # Determine which backend will be used for judge calls
        judge_backend = (
            "nim"       if _NIM_API_KEY and not _API_KEY else
            "anthropic" if _API_KEY else
            "behavioral"
        )

        # Score (P(True) if API key available, else behavioral only)
        if _API_KEY or _NIM_API_KEY:
            result = guard.score_with_ptrue(question, steps, final_answer)
        else:
            result = guard.score_chain(question, steps, final_answer)

        beh_score   = result.behavioral_score
        risk_score  = result.risk_score
        tier        = result.confidence_tier
        ptrue_score = result.behavioral_components.get("ptrue_risk", 0.0) or 0.0

        # RL re-score: if a trained RL model exists, use it as the final risk score.
        # The RL model is trained from human feedback labels — it directly learns what
        # patterns correlate with actual failures in this deployment context.
        rl_score = _rl.predict(risk_score, beh_score, ptrue_score, len(steps), tier)

        # Determine final risk: RL score takes over once trained (RLHF effect)
        final_risk = rl_score if rl_score is not None else risk_score
        final_tier = _risk_to_tier(final_risk)

        # Resolve domain: use provided value or auto-classify from the question text
        domain = args.get("domain") or _classify_domain(question)

        # Re-compute tier using domain-specific thresholds (tighter for high-risk domains)
        final_tier = _risk_to_tier(final_risk, domain)

        # Log to tracker (store judge backend + original behavioral risk for audit trail)
        record = ChainRecord(
            question=question, final_answer=final_answer,
            n_steps=len(steps), risk_score=final_risk, tier=final_tier,
            ptrue_score=ptrue_score, beh_score=beh_score,
            aborted=False, session_id=session_id, source=source,
            metadata={"judge_backend": judge_backend, "base_risk": round(risk_score, 4), "domain": domain},
        )
        chain_id = _tracker.log_chain(record)

        # Auto-trigger retrain check
        should, reason = _rl.should_retrain()
        retrain_msg = None
        if should:
            retrain_result = _rl.retrain(guard)
            retrain_msg = f"RL model auto-retrained: CV AUROC={retrain_result.get('cv_auroc')}"

        response = {
            "chain_id":      chain_id,
            "risk_score":    round(final_risk, 4),
            "tier":          final_tier,
            "needs_review":  final_tier in ("HIGH", "MEDIUM"),
            "needs_alert":   final_tier == "HIGH",
            "beh_score":     round(beh_score, 4),
            "ptrue_score":   round(ptrue_score, 4),
            "rl_score":      round(rl_score, 4) if rl_score is not None else None,
            "rl_active":     rl_score is not None,
            "base_risk":     round(risk_score, 4),   # pre-RL behavioral score (audit)
            "domain":        domain,
            "domain_threshold_adjusted": domain in _DOMAIN_HIGH_THRESHOLDS,
            "interpretation": _interpret(final_risk, final_tier),
            "action":        _recommend_action(final_tier),
        }
        if retrain_msg:
            response["rl_retrain"] = retrain_msg
        return response

    # ── stream_check ───────────────────────────────────────────────────────────
    elif name == "stream_check":
        guard      = _get_guard()
        question   = args["question"]
        steps      = args.get("steps_so_far", [])
        session_id = args.get("session_id", "")

        sg = guard.stream_guard(question, steps, abort_threshold=0.65)

        if sg.abort:
            record = ChainRecord(
                question=question, final_answer="[ABORTED]",
                n_steps=len(steps), risk_score=sg.risk_at_step, tier="HIGH",
                aborted=True, session_id=session_id, source="mcp",
            )
            chain_id = _tracker.log_chain(record)
        else:
            chain_id = None

        return {
            "should_abort": sg.abort,
            "risk_score":   round(sg.risk_at_step, 4),
            "step":         len(steps),
            "chain_id":     chain_id,
            "message": (
                "ABORT: Agent is not finding useful information. Stop and report failure."
                if sg.abort else
                "CONTINUE: Chain looks on-track so far."
            ),
        }

    # ── submit_feedback ────────────────────────────────────────────────────────
    elif name == "submit_feedback":
        chain_id = args["chain_id"]
        correct  = args["correct"]
        note     = args.get("note", "")

        found = _tracker.add_feedback(chain_id, correct)
        if not found:
            return {"error": f"chain_id={chain_id} not found"}

        labeled = _tracker.get_labeled()
        should, reason = _rl.should_retrain()

        response = {
            "status":        "feedback_recorded",
            "chain_id":      chain_id,
            "label":         "correct" if correct else "wrong",
            "n_labeled":     len(labeled),
            "retrain_ready": should,
            "retrain_reason": reason,
        }
        if note:
            response["note"] = note
        if should:
            response["suggestion"] = "Call trigger_retrain to update the RL model with new labels."
        return response

    # ── get_metrics ────────────────────────────────────────────────────────────
    elif name == "get_metrics":
        days    = float(args.get("days", 7))
        summary = _tracker.summary(days=days)
        recent  = _tracker.get_recent(5)

        return {
            "summary": summary,
            "recent_chains": [
                {
                    "id": r["id"],
                    "question": r["question"][:60] + "...",
                    "risk": round(r["risk_score"], 3),
                    "tier": r["tier"],
                    "labeled": r["human_label"] is not None,
                }
                for r in recent
            ],
            "rl_status": {
                "model_trained": _rl._clf is not None,
                "labels_since_last_train": len(_tracker.get_labeled()) - _rl._n_at_last_train,
            },
        }

    # ── get_auroc ──────────────────────────────────────────────────────────────
    elif name == "get_auroc":
        days  = float(args.get("days", 30))
        auroc = _auroc.compute(days=days)
        drift = _auroc.drift_alert()
        hist  = _auroc.history(5)

        return {
            "auroc":           auroc,
            "window_days":     days,
            "n_labeled":       len(_tracker.get_labeled(days=days)),
            "drift_alert":     drift,
            "auroc_history":   [
                {"ts": h["ts"], "auroc": h["auroc"], "n": h["n_labeled"]}
                for h in hist
            ],
            "baseline_ref": {
                "behavioral_only":   0.682,
                "ptrue_ensemble":    0.775,
                "integration_target": 0.795,
            },
            "status": (
                "INSUFFICIENT DATA (need ≥10 labels)" if auroc is None else
                f"AUROC={auroc:.3f} — " + (
                    "EXCELLENT (above integration target)"  if auroc >= 0.795 else
                    "GOOD (above P(True) baseline)"         if auroc >= 0.775 else
                    "ACCEPTABLE (above behavioral baseline)" if auroc >= 0.682 else
                    "BELOW BASELINE — review scoring setup"
                )
            ),
        }

    # ── trigger_retrain ────────────────────────────────────────────────────────
    elif name == "trigger_retrain":
        should, reason = _rl.should_retrain()
        if not should:
            return {"status": "skipped", "reason": reason}
        guard  = _get_guard()
        result = _rl.retrain(guard)
        return result

    # ── get_pending_review ─────────────────────────────────────────────────────
    elif name == "get_pending_review":
        limit   = int(args.get("limit", 10))
        pending = _tracker.pending_feedback(limit=limit)
        labeled = _tracker.get_labeled()
        should, reason = _rl.should_retrain()

        return {
            "pending_count": len(pending),
            "pending": [
                {
                    "chain_id":    r["id"],
                    "question":    r["question"][:80],
                    "answer":      r["final_answer"][:80],
                    "risk_score":  round(r["risk_score"], 3),
                    "tier":        r["tier"],
                    "how_to_label": f"Call submit_feedback with chain_id={r['id']}, correct=true/false",
                }
                for r in pending
            ],
            "total_labeled":  len(labeled),
            "retrain_ready":  should,
            "retrain_reason": reason,
        }

    # ── get_rl_status ──────────────────────────────────────────────────────────
    elif name == "get_rl_status":
        hist       = _rl.checkpoint_history(10)
        auroc_hist = _auroc.history(10)
        labeled    = _tracker.get_labeled()
        should, reason = _rl.should_retrain()

        return {
            "model_trained":   _rl._clf is not None,
            "total_labels":    len(labeled),
            "labels_wrong":    sum(1 for r in labeled if r["human_label"] == 1),
            "labels_correct":  sum(1 for r in labeled if r["human_label"] == 0),
            "labels_since_last_train": len(labeled) - _rl._n_at_last_train,
            "retrain_ready":   should,
            "retrain_reason":  reason,
            "training_history": [
                {
                    "n_labels":  h["n_labels"],
                    "auroc":     h["auroc"],
                    "precision": h["precision"],
                    "recall":    h["recall"],
                    "ts":        h["ts"],
                }
                for h in hist
            ],
            "auroc_trend": [
                {"auroc": h["auroc"], "n": h["n_labeled"], "ts": h["ts"]}
                for h in auroc_hist
            ],
            "drift_alert": _auroc.drift_alert(),
            "rl_loop_config": {
                "retrain_every_n_labels": _rl.RETRAIN_EVERY,
                "min_labels_for_first_train": _rl.MIN_LABELS,
                "min_per_class": _rl.MIN_PER_CLASS,
                "features": ["risk_score", "beh_score", "ptrue_score", "n_steps_norm", "tier_num"],
                "note": "RL learns from human labels regardless of which backend scored the chain. "
                        "Use compare_backends to see per-backend AUROC and pick the best one.",
            },
            "active_backend": (
                "nim"       if _NIM_API_KEY and not _API_KEY else
                "anthropic" if _API_KEY else
                "behavioral"
            ),
        }

    # ── compare_backends ───────────────────────────────────────────────────────
    elif name == "compare_backends":
        days = float(args.get("days", 30))
        comparison = _auroc.compare_backends(days=days)
        summary    = _tracker.backend_summary(days=days)
        active = (
            "nim"       if _NIM_API_KEY and not _API_KEY else
            "anthropic" if _API_KEY else
            "behavioral"
        )
        return {
            "active_backend":   active,
            "days":             days,
            "auroc_by_backend": comparison["backends"],
            "volume_by_backend": summary["backends"],
            "baseline_auroc": {
                "behavioral_only":  0.682,
                "ptrue_haiku":      0.739,
                "ptrue_ensemble":   0.775,
                "integration_target": 0.795,
            },
            "interpretation": (
                "auroc_by_backend shows per-backend scoring quality validated by human labels. "
                "Higher AUROC = that backend better predicts actual failures. "
                "Switch active backend (ANTHROPIC_API_KEY vs NIM_API_KEY in .mcp.json) "
                "to the one with higher AUROC once you have ≥10 labeled chains per backend. "
                "The RL loop then adapts to whichever backend is active."
            ),
        }

    # ── get_nines_dashboard ────────────────────────────────────────────────────
    elif name == "get_nines_dashboard":
        days = float(args.get("days", 30))
        dashboard = _auroc.nines_dashboard(days=days)
        return dashboard

    # ── get_domain_stats ───────────────────────────────────────────────────────
    elif name == "get_domain_stats":
        days = float(args.get("days", 30))
        stats = _tracker.domain_stats(days=days)
        return stats

    # ── check_retry_budget ─────────────────────────────────────────────────────
    elif name == "check_retry_budget":
        attempt_risks  = args.get("attempt_risks", [])
        max_retries    = int(args.get("max_retries", 3))
        low_threshold  = float(args.get("low_threshold", 0.40))
        high_threshold = float(args.get("high_threshold", 0.65))
        result = AgentGuard.should_retry(
            attempt_risks=attempt_risks,
            max_retries=max_retries,
            low_threshold=low_threshold,
            high_threshold=high_threshold,
        )
        return result

    # ── heal_outcome ───────────────────────────────────────────────────────────
    elif name == "heal_outcome":
        chain_id     = int(args["chain_id"])
        success      = bool(args["success"])
        failure_mode = str(args.get("failure_mode", ""))
        action_type  = str(args.get("action_type", ""))
        note         = str(args.get("note", ""))
        row_id = _tracker.log_heal_outcome(
            chain_id=chain_id,
            success=success,
            failure_mode=failure_mode,
            action_type=action_type,
            note=note,
        )
        return {
            "heal_outcome_id": row_id,
            "chain_id":        chain_id,
            "recorded":        True,
            "success":         success,
            "failure_mode":    failure_mode or "unspecified",
            "action_type":     action_type  or "unspecified",
            "note":            "Call get_heal_stats to see aggregated repair effectiveness.",
        }

    # ── get_heal_stats ─────────────────────────────────────────────────────────
    elif name == "get_heal_stats":
        days  = float(args.get("days", 30))
        stats = _tracker.heal_outcome_stats(days=days)
        if stats["n_outcomes"] == 0:
            return {
                **stats,
                "interpretation": (
                    "No heal outcomes recorded yet. "
                    "Call heal_outcome after each SelfHealer repair attempt to build this dataset."
                ),
            }
        lines = [
            f"Overall success rate: {stats['success_rate']:.1%} over {stats['n_outcomes']} repairs."
        ]
        for mode, v in stats["by_mode"].items():
            lines.append(
                f"  {mode}: {v['n_success']}/{v['n']} succeeded "
                f"({v['success_rate']:.1%})"
            )
        return {
            **stats,
            "interpretation": " ".join(lines),
        }

    # ── retrieval_audit ────────────────────────────────────────────────────────
    elif name == "retrieval_audit":
        question = args.get("question", "")
        steps    = args.get("steps", [])
        guard    = _get_guard()
        rq       = guard.retrieval_quality(question, steps)

        # Annotate each per-step similarity with the step's search query
        search_steps = [s for s in steps if s.get("action_type") == "Search"
                        and s.get("observation", "").strip()]
        step_details = []
        for i, (sim, step) in enumerate(zip(rq.get("per_step", []), search_steps)):
            step_details.append({
                "step_index":  i,
                "query":       step.get("action_arg", "")[:120],
                "cosine_sim":  round(float(sim), 4),
                "relevant":    float(sim) > 0.25,
            })

        return {
            "quality_label":    rq["quality_label"],
            "mean_sim":         rq["mean_sim"],
            "max_sim":          rq["max_sim"],
            "min_sim":          rq["min_sim"],
            "var_sim":          rq["var_sim"],
            "n_relevant_steps": rq["n_relevant"],
            "n_search_steps":   len(search_steps),
            "per_step_details": step_details,
            "interpretation": (
                f"Retrieval quality is {rq['quality_label']}. "
                f"{rq['n_relevant']}/{len(search_steps)} steps returned relevant observations "
                f"(cosine_sim > 0.25). "
                + ("Consider rephrasing search queries to be more specific. "
                   if rq["mean_sim"] < 0.35 else
                   "At least some observations are topically relevant. "
                   if rq["mean_sim"] < 0.50 else
                   "Retrieval looks healthy — failure likely in reasoning, not search.")
            ),
        }

    else:
        return {"error": f"Unknown tool: {name}"}


# ── Helpers ────────────────────────────────────────────────────────────────────

def _classify_domain(question: str) -> str:
    """Keyword-based domain classifier. Returns one of 7 domain tags."""
    q = question.lower()
    # Person/biographical
    if any(w in q for w in ("who is", "who was", "who invented", "who founded", "who wrote",
                             "who played", "who won", "who led", "who created", "biography")):
        return "person_factual"
    # Temporal
    if any(w in q for w in ("when did", "when was", "what year", "what date", "in what year",
                             "how long ago", "since when", "what century", "timeline")):
        return "temporal_factual"
    # Location/geography
    if any(w in q for w in ("where is", "where was", "where did", "what country", "what city",
                             "what state", "located in", "geography", "capital of")):
        return "location_factual"
    # Numeric/calculation (word-boundary aware: "count " not "country")
    if any(w in q for w in ("how many", "how much", "calculate", "compute", "what is the number",
                             "total", "count of", "count the", "percentage", "how far", "distance")):
        return "numeric"
    # Reasoning/causal
    if any(w in q for w in ("why did", "why is", "explain", "what causes", "reason for",
                             "how does", "what led to", "effect of", "impact of", "consequence")):
        return "reasoning"
    # Generative/creative (less suited to ReAct search)
    if any(w in q for w in ("write", "generate", "create", "draft", "compose", "list all",
                             "summarize", "describe", "tell me about")):
        return "generative"
    return "general"


# Domain-specific HIGH thresholds (empirically derived from real TriviaQA analysis).
# person_factual/temporal/numeric fail at 2× the global rate, so alert earlier.
_DOMAIN_HIGH_THRESHOLDS = {
    "person_factual":   0.55,   # 31% error → lower bar to alert
    "temporal_factual": 0.50,   # 50% error (small sample) → be aggressive
    "numeric":          0.55,   # 25% error → slightly tighter
    # general/location/reasoning/generative: use default 0.65
}
_DOMAIN_MEDIUM_THRESHOLDS = {
    "person_factual":   0.35,
    "temporal_factual": 0.30,
    "numeric":          0.35,
}


def _risk_to_tier(risk: float, domain: str = "general") -> str:
    """Convert risk score to tier, using tighter thresholds for high-risk domains."""
    high_t   = _DOMAIN_HIGH_THRESHOLDS.get(domain, 0.65)
    medium_t = _DOMAIN_MEDIUM_THRESHOLDS.get(domain, 0.40)
    if risk >= high_t:
        return "HIGH"
    elif risk >= medium_t:
        return "MEDIUM"
    else:
        return "LOW"


def _interpret(risk: float, tier: str) -> str:
    if tier == "LOW":
        return f"LOW risk ({risk:.3f}): Agent found clear evidence. Answer is likely reliable."
    elif tier == "MEDIUM":
        return f"MEDIUM risk ({risk:.3f}): Some uncertainty detected. Spot-check recommended."
    else:
        return f"HIGH risk ({risk:.3f}): Agent failed, contradicted itself, or gave up. Do not use without verification."


def _recommend_action(tier: str) -> str:
    return {
        "LOW":    "Use answer directly. No review needed.",
        "MEDIUM": "Show answer to user with caveat; log for periodic review.",
        "HIGH":   "Block answer. Alert human reviewer. Do not send to end user.",
    }.get(tier, "Review manually.")


# ── Entry point ────────────────────────────────────────────────────────────────

async def main():
    import argparse
    parser = argparse.ArgumentParser(description="llm-guard-kit MCP Server")
    parser.add_argument("--sse", action="store_true", help="Use SSE transport instead of stdio")
    parser.add_argument("--port", type=int, default=8765, help="Port for SSE transport")
    args = parser.parse_args()

    if args.sse:
        from mcp.server.sse import SseServerTransport
        from starlette.applications import Starlette
        from starlette.routing import Route
        import uvicorn

        sse = SseServerTransport("/messages/")

        async def handle_sse(request):
            async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
                await app.run(streams[0], streams[1], app.create_initialization_options())

        starlette_app = Starlette(routes=[Route("/sse", endpoint=handle_sse)])
        print(f"llm-guard-kit MCP server running on http://localhost:{args.port}/sse", file=sys.stderr)
        uvicorn.run(starlette_app, host="0.0.0.0", port=args.port)
    else:
        print("llm-guard-kit MCP server starting (stdio transport)...", file=sys.stderr)
        async with stdio_server() as (read_stream, write_stream):
            await app.run(read_stream, write_stream, app.create_initialization_options())


def run():
    """Sync entry point for `llm-guard-mcp` CLI script (registered in pyproject.toml)."""
    asyncio.run(main())


if __name__ == "__main__":
    run()
