"""
QPPG HTTP API Server — FastAPI wrapper around QPPGService.

v0.3.0 additions:
  - /dashboard and /dashboard/api/* — monitoring UI (plain HTML, no extra deps)
  - /api/keys          — create/manage API keys for SaaS auth
  - /api/{domain}/score — authenticated scoring endpoint
  - /api/{domain}/audit-log — export audit log
  - Per-domain SQLite persistence via ChainStore
  - In-memory rate limiting (1000 req/hr per API key)

Original endpoints:
  POST /score          — score a reasoning chain
  POST /calibrate      — add a verified chain to the calibration pool
  POST /bulk-calibrate — add many chains at once from existing logs
  GET  /status         — deployment readiness and stats
  GET  /progress       — ASCII progress bar
  POST /reset          — clear calibration pool

Run:
    llm-guard-kit serve --domain my_domain --port 8000
    # or:
    python -m qppg_service.server --domain my_domain --port 8000

Dashboard:
    llm-guard-kit dashboard --domain my_domain --port 8080
    # Then open http://localhost:8080/dashboard
"""

from __future__ import annotations

import argparse
import collections
import hashlib
import os
import sys
import time
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from qppg_service.service import QPPGService


# ── Request / Response models ─────────────────────────────────────────────────

class Step(BaseModel):
    thought:     str = ""
    action_type: str = "Finish"
    action_arg:  str = ""
    observation: str = ""


class ScoreRequest(BaseModel):
    question: str
    steps:    list[Step] = []
    raw_text: str = ""


class CalibrateRequest(BaseModel):
    question: str
    steps:    list[Step] = []
    correct:  bool
    raw_text: str = ""
    metadata: dict = {}


class BulkCalibrateRequest(BaseModel):
    chains: list[CalibrateRequest]


class CreateKeyRequest(BaseModel):
    customer_id:   str
    domain_prefix: str = ""


class UserCreateKeyRequest(BaseModel):
    domain_prefix: str = ""
    label:         str = ""


class StreamScoreRequest(BaseModel):
    """Request body for POST /score/stream (SSE step-by-step scoring)."""
    question:  str
    steps:     list[Step]
    threshold: float = 0.70


class ImportGmmRequest(BaseModel):
    """Request body for POST /api/{domain}/import-gmm (federated GMM)."""
    params:   dict
    n_chains: int = 0


class WorkflowSaveRequest(BaseModel):
    name:  str = "Untitled workflow"
    nodes: list = []
    edges: list = []
    id:    Optional[int] = None


class GuardrailSaveRequest(BaseModel):
    domain:        str = ""
    threshold:     float = 0.65
    action:        str   = "alert"
    webhook_url:   str   = ""
    slack_webhook: str   = ""
    enabled:       bool  = True


class PipelineSaveRequest(BaseModel):
    name:  str = "Untitled pipeline"
    nodes: list = []
    edges: list = []
    id:    Optional[int] = None


class PipelineRunRequest(BaseModel):
    query: str


class ABCompareRequest(BaseModel):
    pipeline_a_id: int
    pipeline_b_id: int
    query: str


class TestSetSaveRequest(BaseModel):
    name:        str = "Untitled test set"
    cases:       list = []   # [{query, expected_output, risk_threshold}]
    pipeline_id: Optional[int] = None
    id:          Optional[int] = None


class CalibrateFitRequest(BaseModel):
    """Request body for POST /v2/calibrate/fit."""
    labeled_runs: list[dict]        # [{question, steps, correct}, ...]
    ptrue_values: list[float] = []  # optional P(True) per chain


class TestRunRequest(BaseModel):
    pipeline_id: int
    notify_email: str = ""


class OrgCreateRequest(BaseModel):
    name: str


class OrgInviteRequest(BaseModel):
    email: str


class OrgAcceptRequest(BaseModel):
    token: str


class PipelineVersionNoteRequest(BaseModel):
    note: str = ""


class EvaluateRequest(BaseModel):
    text:       str
    context:    str = ""
    evaluators: list[str] = []   # empty = run all


class TestABRequest(BaseModel):
    test_set_id:   int
    pipeline_a_id: int
    pipeline_b_id: int
    notify_email:  str = ""


# ── Rate limiter (in-memory, per API key) ─────────────────────────────────────

class _RateLimiter:
    """Simple sliding-window rate limiter; max_per_hour requests per key."""

    def __init__(self, max_per_hour: int = 1000):
        self._max    = max_per_hour
        self._counts: dict[str, list[float]] = collections.defaultdict(list)

    def check(self, key_hash: str) -> bool:
        now  = time.time()
        cutoff = now - 3600
        q = [t for t in self._counts[key_hash] if t > cutoff]
        self._counts[key_hash] = q
        if len(q) >= self._max:
            return False
        self._counts[key_hash].append(now)
        return True


_rate_limiter = _RateLimiter()


# ── App factory ───────────────────────────────────────────────────────────────

def create_app(
    service: QPPGService,
    db_path: Optional[str] = None,
    admin_key: Optional[str] = None,
    enable_dashboard: bool = False,
    dashboard_key: Optional[str] = None,
) -> FastAPI:
    """
    Build the FastAPI application.

    Parameters
    ----------
    service          : QPPGService for the original /score, /calibrate endpoints
    db_path          : SQLite database for SaaS auth + dashboard (None = disabled)
    admin_key        : secret string required in X-Admin-Key header to create API keys
    enable_dashboard : mount /dashboard and /dashboard/api/* routes
    dashboard_key    : if set, all /dashboard/api/admin/* routes require
                       ``X-Dashboard-Key: <value>`` header
    """
    app = FastAPI(
        title="llm-guard-kit API",
        description=(
            "Real-time confidence scoring and monitoring for LLM agent chains. "
            "v0.5.0 — SaaS auth, Stripe billing, Free/Pro plans, Next.js frontend."
        ),
        version="0.5.0",
    )

    # CORS — required for Vercel frontend → Railway backend cross-origin requests
    app.add_middleware(
        CORSMiddleware,
        allow_origins=os.environ.get("CORS_ORIGINS", "http://localhost:3000").split(","),
        allow_methods=["*"],
        allow_headers=["*"],
        allow_credentials=True,
    )

    # Optional SQLite store
    store = None
    if db_path:
        from .store import ChainStore
        store = ChainStore(db_path)

    # Mount SaaS auth + billing routers (requires store)
    if store is not None:
        from .auth import make_auth_router, _get_current_user_factory
        from .billing import make_billing_router
        _user_auth_dep = _get_current_user_factory(store)
        app.include_router(make_auth_router(store), prefix="/auth", tags=["auth"])
        app.include_router(make_billing_router(store, _user_auth_dep), prefix="/billing", tags=["billing"])

    # Per-domain LabelFreeScorer cache (for federated GMM endpoints)
    from .label_free_scorer import LabelFreeScorer as _LFS
    _domain_scorers: dict[str, _LFS] = {}

    # ── Original endpoints (backward-compatible) ───────────────────────────────

    @app.post("/score", summary="Score a reasoning chain")
    async def score(req: ScoreRequest):
        """
        Returns a confidence score for an agent reasoning chain.
        - **confidence**: 0–1 (None during cold start)
        - **needs_review**: True if below threshold or cold start
        """
        steps  = [s.model_dump() for s in req.steps]
        result = service.score(req.question, steps=steps, raw_text=req.raw_text)
        return result.to_dict()

    @app.post("/calibrate", summary="Add a verified chain to the calibration pool")
    async def calibrate(req: CalibrateRequest):
        steps  = [s.model_dump() for s in req.steps]
        result = service.add_verified(
            req.question, steps=steps,
            correct=req.correct, raw_text=req.raw_text, metadata=req.metadata,
        )
        return result.to_dict()

    @app.post("/bulk-calibrate", summary="Add many verified chains at once")
    async def bulk_calibrate(req: BulkCalibrateRequest):
        chains = [{
            "question": c.question,
            "steps":    [s.model_dump() for s in c.steps],
            "correct":  c.correct,
            "raw_text": c.raw_text,
            "metadata": c.metadata,
        } for c in req.chains]
        return service.bulk_calibrate(chains)

    @app.post("/v2/calibrate/fit", summary="Fit a DeepLocalVerifier on labeled chains (v0.16.0)")
    async def calibrate_fit(req: CalibrateFitRequest):
        """
        Fit a DeepLocalVerifier on labeled chains and return the serialized model.

        Input: list of {question, steps, correct} labeled chains.
        Returns: base64-encoded pickle of the fitted model + 5-fold CV AUROC.

        Security: model_b64 is a pickle — only share with trusted clients.
        Requires: pip install llm-guard-kit[qara]
        """
        from qppg_service.calibration_service import fit_verifier
        try:
            result = fit_verifier(
                req.labeled_runs,
                ptrue_values=req.ptrue_values if req.ptrue_values else None,
            )
            return result
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc))
        except RuntimeError as exc:
            raise HTTPException(status_code=422, detail=str(exc))

    @app.get("/status", summary="Deployment readiness and statistics")
    async def status():
        return service.status()

    @app.get("/progress", summary="ASCII progress bar")
    async def progress():
        return {"progress": service.progress_bar()}

    @app.post("/reset", summary="Clear calibration pool — DESTRUCTIVE")
    async def reset(confirm: str = ""):
        if confirm != "YES_RESET":
            raise HTTPException(400, "Pass query param confirm=YES_RESET to confirm.")
        service.reset()
        return {"reset": True, "domain": service.domain_name}

    # ── SSE streaming: score step-by-step as the agent runs ───────────────────
    #
    # Client sends steps one at a time via Server-Sent Events.
    # The server scores each prefix and streams back intermediate risk scores,
    # enabling early termination after t=2 steps saves ~74% of tokens (exp76).
    #
    # Protocol:
    #   POST /score/stream  — opens an SSE stream
    #   Body: { "question": "...", "steps": [...], "threshold": 0.70 }
    #   The server scores prefix(1), prefix(2), ..., prefix(N) and emits:
    #     data: {"t": 1, "risk_score": 0.42, "stop": false}
    #     data: {"t": 2, "risk_score": 0.71, "stop": true}   ← threshold exceeded
    #   After all steps or early stop, emits:
    #     data: {"t": N, "risk_score": 0.71, "stop": true, "done": true}

    from fastapi.responses import StreamingResponse as _StreamingResp
    import json as _json
    import asyncio as _asyncio

    @app.post(
        "/score/stream",
        summary="Score steps via SSE as agent runs (early stopping, exp76 t=2)",
        response_description="Server-Sent Events stream of per-step risk scores",
    )
    async def score_stream(req: StreamScoreRequest):
        """
        Stream intermediate risk scores for each step prefix (exp76).

        Returns ``text/event-stream`` (SSE). Each event carries:
        - ``t``           : step index (1-indexed)
        - ``risk_score``  : cumulative behavioral risk score for first t steps
        - ``stop``        : True if risk_score ≥ threshold → early terminate
        - ``done``        : True on the final event (all steps scored or stopped)
        - ``failure_mode``: detected failure mode (on final event)

        **Early stopping**: optimal at t=2 with ~74% token savings (exp76).
        Set threshold=0.70 (default) to match exp67 recommendation.
        """
        from .label_free_scorer import LabelFreeScorer as _LFS

        steps   = [s.model_dump() for s in req.steps]
        scorer  = service._scorer if hasattr(service, "_scorer") else _LFS()
        question = req.question
        threshold = req.threshold

        async def _event_gen():
            stopped = False
            for t in range(1, len(steps) + 1):
                prefix = steps[:t]
                is_last = (t == len(steps))
                # Score only behavioral (fast, no calibration) for streaming
                # Use the service scorer if available (may be calibrated)
                try:
                    result = scorer.score(
                        question, prefix,
                        final_answer="" if not is_last else steps[-1].get("action_arg", ""),
                        finished=is_last,
                    )
                except Exception:
                    result = None

                risk = round(result.risk_score, 4) if result else 0.5
                stop = risk >= threshold
                payload: dict = {"t": t, "risk_score": risk, "stop": stop or stopped}
                if is_last or stop:
                    payload["done"] = True
                    if result:
                        payload["failure_mode"] = result.failure_mode
                yield "data: " + _json.dumps(payload) + "\n\n"
                if stop:
                    stopped = True
                    break
                await _asyncio.sleep(0)  # yield control to event loop

        return _StreamingResp(
            _event_gen(),
            media_type="text/event-stream",
            headers={
                "Cache-Control":   "no-cache",
                "X-Accel-Buffering": "no",
            },
        )

    # ── SaaS: API key management ───────────────────────────────────────────────

    if store is not None:

        def _require_admin(x_admin_key: Optional[str]):
            if admin_key and x_admin_key != admin_key:
                raise HTTPException(401, "Invalid or missing X-Admin-Key header.")

        def _auth_customer(authorization: Optional[str]) -> dict:
            if not authorization or not authorization.startswith("Bearer "):
                raise HTTPException(401, "Authorization: Bearer <key> required.")
            raw_key  = authorization.removeprefix("Bearer ").strip()
            key_info = store.verify_api_key(raw_key)
            if not key_info:
                raise HTTPException(401, "Invalid or inactive API key.")
            key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
            if not _rate_limiter.check(key_hash):
                raise HTTPException(429, "Rate limit exceeded (1000 req/hr).")
            return key_info

        @app.post("/api/keys", summary="Create an API key (admin only)")
        async def create_key(
            req:         CreateKeyRequest,
            x_admin_key: Optional[str] = Header(None),
        ):
            _require_admin(x_admin_key)
            raw_key = store.create_api_key(req.customer_id, req.domain_prefix)
            return {
                "api_key":       raw_key,
                "key_id":        hashlib.sha256(raw_key.encode()).hexdigest()[:12],
                "customer_id":   req.customer_id,
                "domain_prefix": req.domain_prefix,
                "note":          "Save this key — it is not stored in plaintext.",
            }

        @app.get("/api/keys/{key_id}/status", summary="API key usage stats")
        async def key_status(key_id: str):
            try:
                row = store.get_key_by_prefix(key_id)
            except ValueError as e:
                raise HTTPException(409, str(e))
            if not row:
                raise HTTPException(404, "Key not found.")
            return row

        # ── User self-serve API key management (JWT auth) ──────────────────────

        from fastapi import Depends as _Depends
        from .auth import _get_current_user_factory as _ucf
        _jwt_user = _ucf(store)

        @app.post("/user/keys", summary="Create an API key (user self-serve)")
        async def user_create_key(
            req:  UserCreateKeyRequest,
            user: dict = _Depends(_jwt_user),
        ):
            raw_key = store.create_api_key(
                customer_id=f"user_{user['id']}",
                domain_prefix=req.domain_prefix,
            )
            key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
            store.link_api_key_to_user(key_hash, user["id"])
            return {
                "api_key":       raw_key,
                "key_prefix":    key_hash[:12],
                "customer_id":   f"user_{user['id']}",
                "domain_prefix": req.domain_prefix,
                "note":          "Save this key — it is not stored in plaintext.",
            }

        @app.get("/user/keys", summary="List API keys for the current user")
        async def user_list_keys(user: dict = _Depends(_jwt_user)):
            return store.get_keys_for_user(user["id"])

        @app.delete("/user/keys/{key_prefix}", summary="Revoke a user API key")
        async def user_revoke_key(
            key_prefix: str,
            user:       dict = _Depends(_jwt_user),
        ):
            keys = store.get_keys_for_user(user["id"])
            if not any(k["key_prefix"] == key_prefix for k in keys):
                raise HTTPException(403, "Key not found or does not belong to you.")
            revoked = store.revoke_key_by_prefix(key_prefix)
            if not revoked:
                raise HTTPException(404, "Key not found.")
            return {"revoked": True, "key_prefix": key_prefix}

        # ── Domain stats ──────────────────────────────────────────────────────

        @app.get("/user/domains", summary="List domains + stats for this user")
        async def user_domains(user: dict = _Depends(_jwt_user)):
            return store.get_domains_for_user(user["id"])

        @app.get("/user/chains/recent", summary="Recent scored chains for this user")
        async def user_recent_chains(
            limit: int = 20,
            user:  dict = _Depends(_jwt_user),
        ):
            return store.get_recent_chains_for_user(user["id"], limit=min(limit, 100))

        # ── Workflow management ───────────────────────────────────────────────

        @app.get("/user/workflows", summary="List workflows")
        async def list_workflows(user: dict = _Depends(_jwt_user)):
            return store.get_workflows(user["id"])

        @app.get("/user/workflows/{workflow_id}", summary="Get a single workflow")
        async def get_workflow(workflow_id: int, user: dict = _Depends(_jwt_user)):
            w = store.get_workflow(workflow_id, user["id"])
            if not w:
                raise HTTPException(404, "Workflow not found.")
            return w

        @app.post("/user/workflows", summary="Save a workflow")
        async def save_workflow(req: WorkflowSaveRequest, user: dict = _Depends(_jwt_user)):
            return store.save_workflow(user["id"], req.name, req.nodes, req.edges, req.id)

        @app.patch("/user/workflows/{workflow_id}/toggle", summary="Enable or disable a workflow")
        async def toggle_workflow(workflow_id: int, body: dict, user: dict = _Depends(_jwt_user)):
            w = store.toggle_workflow(workflow_id, user["id"], bool(body.get("enabled", True)))
            if not w:
                raise HTTPException(404, "Workflow not found.")
            return w

        @app.delete("/user/workflows/{workflow_id}", summary="Delete a workflow")
        async def delete_workflow(workflow_id: int, user: dict = _Depends(_jwt_user)):
            ok = store.delete_workflow(workflow_id, user["id"])
            if not ok:
                raise HTTPException(404, "Workflow not found.")
            return {"deleted": True}

        # ── Guardrail management ──────────────────────────────────────────────

        @app.get("/user/guardrails", summary="List guardrails for this user")
        async def list_guardrails(user: dict = _Depends(_jwt_user)):
            return store.get_guardrails(user["id"])

        @app.post("/user/guardrails", summary="Save a guardrail rule")
        async def save_guardrail(req: GuardrailSaveRequest, user: dict = _Depends(_jwt_user)):
            return store.save_guardrail(
                user["id"], req.domain, req.threshold,
                req.action, req.webhook_url, req.slack_webhook, req.enabled,
            )

        @app.delete("/user/guardrails/{domain}", summary="Delete a guardrail")
        async def delete_guardrail(domain: str, user: dict = _Depends(_jwt_user)):
            ok = store.delete_guardrail(user["id"], domain)
            if not ok:
                raise HTTPException(404, "Guardrail not found.")
            return {"deleted": True}

        # ── Pipeline management ───────────────────────────────────────────────

        @app.get("/user/pipelines", summary="List pipelines")
        async def list_pipelines(user: dict = _Depends(_jwt_user)):
            return store.get_pipelines(user["id"])

        @app.post("/user/pipelines", summary="Save a pipeline")
        async def save_pipeline(req: PipelineSaveRequest, user: dict = _Depends(_jwt_user)):
            return store.save_pipeline(user["id"], req.name, req.nodes, req.edges, req.id)

        @app.get("/user/pipelines/{pipeline_id}", summary="Get a single pipeline")
        async def get_pipeline(pipeline_id: int, user: dict = _Depends(_jwt_user)):
            p = store.get_pipeline(pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")
            return p

        @app.patch("/user/pipelines/{pipeline_id}/toggle", summary="Enable or disable a pipeline")
        async def toggle_pipeline(pipeline_id: int, body: dict, user: dict = _Depends(_jwt_user)):
            p = store.toggle_pipeline(pipeline_id, user["id"], bool(body.get("enabled", True)))
            if not p:
                raise HTTPException(404, "Pipeline not found.")
            return p

        @app.delete("/user/pipelines/{pipeline_id}", summary="Delete a pipeline")
        async def delete_pipeline(pipeline_id: int, user: dict = _Depends(_jwt_user)):
            ok = store.delete_pipeline(pipeline_id, user["id"])
            if not ok:
                raise HTTPException(404, "Pipeline not found.")
            return {"deleted": True}

        @app.post("/user/pipelines/{pipeline_id}/run", summary="Execute a pipeline (sync)")
        async def run_pipeline_endpoint(
            pipeline_id: int,
            req:         PipelineRunRequest,
            user:        dict = _Depends(_jwt_user),
        ):
            p = store.get_pipeline(pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")

            user_configs = store.get_user_config(user["id"])
            run_rec      = store.create_run(pipeline_id, user["id"], req.query)
            run_id       = run_rec["id"]

            from .pipeline_executor import run_pipeline
            result = await run_pipeline(p, req.query, user["id"], user_configs, store)

            store.update_run(
                run_id,
                result=result.get("result", ""),
                trace=result.get("trace", []),
                risk_score=result.get("risk_score"),
                guardrail_action=result.get("guardrail_action"),
                status=result.get("status", "done"),
            )
            return {**result, "run_id": run_id}

        @app.post("/user/pipelines/{pipeline_id}/run/stream", summary="Execute a pipeline (SSE streaming)")
        async def run_pipeline_stream(
            pipeline_id: int,
            req:         PipelineRunRequest,
            user:        dict = _Depends(_jwt_user),
        ):
            """
            Stream pipeline execution as Server-Sent Events.
            Each event is a JSON object with "type" field:
              node_start  — a node began executing
              node_done   — a node finished (includes output + duration_ms)
              node_error  — a node raised an exception
              done        — pipeline finished (includes result, trace, risk_score)
            """
            p = store.get_pipeline(pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")

            user_configs = store.get_user_config(user["id"])
            run_rec      = store.create_run(pipeline_id, user["id"], req.query)
            run_id       = run_rec["id"]

            from .pipeline_executor import stream_pipeline
            from fastapi.responses import StreamingResponse as _SR

            async def _event_gen():
                final: dict = {}
                async for event in stream_pipeline(p, req.query, user["id"], user_configs, store):
                    if event["type"] == "done":
                        final = event
                    payload = _json.dumps({**event, "run_id": run_id})
                    yield f"data: {payload}\n\n"
                    await _asyncio.sleep(0)  # yield to event loop

                # Persist final result
                store.update_run(
                    run_id,
                    result=final.get("result", ""),
                    trace=final.get("trace", []),
                    risk_score=final.get("risk_score"),
                    guardrail_action=final.get("guardrail_action"),
                    status=final.get("status", "done"),
                )

            return _SR(
                _event_gen(),
                media_type="text/event-stream",
                headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
            )

        @app.get("/user/pipelines/{pipeline_id}/runs", summary="Pipeline run history")
        async def get_pipeline_runs(pipeline_id: int, user: dict = _Depends(_jwt_user)):
            p = store.get_pipeline(pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")
            return store.get_runs(pipeline_id, user["id"])

        # ── User config (LLM API keys) ────────────────────────────────────────

        @app.get("/user/configs", summary="Get user LLM API key config")
        async def get_user_configs(user: dict = _Depends(_jwt_user)):
            cfg = store.get_user_config(user["id"])
            # Mask key values for security — return prefix only
            masked = {}
            for k, v in cfg.items():
                if v and len(v) > 8:
                    masked[k] = v[:8] + "..." + v[-4:]
                else:
                    masked[k] = v
            return masked

        @app.post("/user/configs", summary="Save user LLM API key config")
        async def set_user_configs(body: dict, user: dict = _Depends(_jwt_user)):
            # Merge with existing (don't overwrite keys not provided)
            existing = store.get_user_config(user["id"])
            for k, v in body.items():
                if v:  # only update non-empty values
                    existing[k] = v
            store.set_user_config(user["id"], existing)
            return {"saved": True}

        # ── Time-series analytics ─────────────────────────────────────────────

        @app.get("/user/timeseries", summary="Bucketed risk score time-series")
        async def get_timeseries(
            domain: str = "",
            days:   int = 30,
            bucket: str = "day",
            user:   dict = _Depends(_jwt_user),
        ):
            return store.get_timeseries(
                user["id"],
                domain=domain or None,
                days=min(days, 365),
                bucket=bucket,
            )

        # ── Pipeline versioning ───────────────────────────────────────────────

        @app.get("/user/pipelines/{pipeline_id}/versions", summary="List pipeline versions")
        async def list_versions(pipeline_id: int, user: dict = _Depends(_jwt_user)):
            p = store.get_pipeline(pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")
            return store.get_pipeline_versions(pipeline_id, user["id"])

        @app.post("/user/pipelines/{pipeline_id}/versions", summary="Snapshot pipeline as new version")
        async def create_version(
            pipeline_id: int,
            req: PipelineVersionNoteRequest,
            user: dict = _Depends(_jwt_user),
        ):
            p = store.get_pipeline(pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")
            return store.snapshot_pipeline_version(pipeline_id, user["id"], note=req.note)

        @app.post("/user/pipelines/{pipeline_id}/versions/{version_id}/restore",
                  summary="Restore pipeline to a saved version")
        async def restore_version(
            pipeline_id: int,
            version_id:  int,
            user: dict = _Depends(_jwt_user),
        ):
            p = store.restore_pipeline_version(version_id, pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Version not found or belongs to another pipeline.")
            return p

        @app.get("/user/pipelines/{pipeline_id}/versions/{version_id}",
                 summary="Get a specific pipeline version (with nodes/edges)")
        async def get_version(pipeline_id: int, version_id: int, user: dict = _Depends(_jwt_user)):
            v = store.get_pipeline_version(version_id, user["id"])
            if not v or v["pipeline_id"] != pipeline_id:
                raise HTTPException(404, "Version not found.")
            return v

        @app.get("/user/pipelines/{pipeline_id}/versions/{version_a_id}/diff/{version_b_id}",
                 summary="Diff two pipeline versions — returns added/removed/changed nodes+edges")
        async def diff_versions(
            pipeline_id:  int,
            version_a_id: int,
            version_b_id: int,
            user: dict = _Depends(_jwt_user),
        ):
            va = store.get_pipeline_version(version_a_id, user["id"])
            vb = store.get_pipeline_version(version_b_id, user["id"])
            if not va or va["pipeline_id"] != pipeline_id:
                raise HTTPException(404, f"Version {version_a_id} not found.")
            if not vb or vb["pipeline_id"] != pipeline_id:
                raise HTTPException(404, f"Version {version_b_id} not found.")
            diff = store.diff_pipeline_versions(va, vb)
            return {
                "pipeline_id": pipeline_id,
                "version_a":   {"id": version_a_id, "version": va["version"], "note": va.get("note", "")},
                "version_b":   {"id": version_b_id, "version": vb["version"], "note": vb.get("note", "")},
                **diff,
            }

        # ── A/B comparison ────────────────────────────────────────────────────

        @app.post("/user/compare", summary="Run two pipelines side-by-side (A/B)")
        async def run_ab_compare(req: ABCompareRequest, user: dict = _Depends(_jwt_user)):
            pa = store.get_pipeline(req.pipeline_a_id, user["id"])
            pb = store.get_pipeline(req.pipeline_b_id, user["id"])
            if not pa:
                raise HTTPException(404, f"Pipeline A ({req.pipeline_a_id}) not found.")
            if not pb:
                raise HTTPException(404, f"Pipeline B ({req.pipeline_b_id}) not found.")

            cmp_rec      = store.create_comparison(user["id"], req.pipeline_a_id, req.pipeline_b_id, req.query)
            user_configs = store.get_user_config(user["id"])

            from .pipeline_executor import run_pipeline
            ra, rb = await _asyncio.gather(
                run_pipeline(pa, req.query, user["id"], user_configs, store),
                run_pipeline(pb, req.query, user["id"], user_configs, store),
                return_exceptions=True,
            )

            result_a = ra.get("result", "") if isinstance(ra, dict) else str(ra)
            result_b = rb.get("result", "") if isinstance(rb, dict) else str(rb)
            trace_a  = ra.get("trace", []) if isinstance(ra, dict) else []
            trace_b  = rb.get("trace", []) if isinstance(rb, dict) else []
            risk_a   = ra.get("risk_score") if isinstance(ra, dict) else None
            risk_b   = rb.get("risk_score") if isinstance(rb, dict) else None

            store.update_comparison(
                cmp_rec["id"], result_a, result_b, trace_a, trace_b, risk_a, risk_b
            )
            return store.get_comparison(cmp_rec["id"], user["id"])

        @app.get("/user/compare", summary="List A/B comparisons")
        async def list_comparisons(user: dict = _Depends(_jwt_user)):
            return store.get_comparisons(user["id"])

        @app.get("/user/compare/{cmp_id}", summary="Get a single A/B comparison")
        async def get_comparison(cmp_id: int, user: dict = _Depends(_jwt_user)):
            c = store.get_comparison(cmp_id, user["id"])
            if not c:
                raise HTTPException(404, "Comparison not found.")
            return c

        @app.post("/user/compare/testset", summary="A/B compare two pipelines over a full test set")
        async def run_ab_testset(req: TestABRequest, user: dict = _Depends(_jwt_user)):
            """
            Runs every case in the test set through pipeline A and pipeline B.
            Returns per-case results and aggregate win/tie stats.
            Runs in background — returns {run_a_id, run_b_id} immediately.
            """
            ts = store.get_test_set(req.test_set_id, user["id"])
            if not ts:
                raise HTTPException(404, "Test set not found.")
            pa = store.get_pipeline(req.pipeline_a_id, user["id"])
            if not pa:
                raise HTTPException(404, f"Pipeline A ({req.pipeline_a_id}) not found.")
            pb = store.get_pipeline(req.pipeline_b_id, user["id"])
            if not pb:
                raise HTTPException(404, f"Pipeline B ({req.pipeline_b_id}) not found.")

            run_a = store.create_test_run(req.test_set_id, user["id"], req.pipeline_a_id)
            run_b = store.create_test_run(req.test_set_id, user["id"], req.pipeline_b_id)
            user_configs = store.get_user_config(user["id"])

            from .pipeline_executor import run_pipeline as _run_pipeline

            async def _run_one_pipeline(pipeline, run_id: int) -> None:
                cases        = ts.get("cases", [])
                results:list = []
                for case in cases:
                    query  = case.get("query", "")
                    thresh = float(case.get("risk_threshold", 0.65))
                    try:
                        r = await _run_pipeline(pipeline, query, user["id"], user_configs, store)
                        risk   = r.get("risk_score", 0.0) or 0.0
                        passed = risk < thresh
                        results.append({
                            "query": query, "output": r.get("result", ""),
                            "risk_score": risk, "passed": passed,
                            "status": r.get("status", "done"),
                        })
                    except Exception as e:
                        results.append({
                            "query": query, "output": "", "risk_score": None,
                            "passed": False, "status": "error", "error": str(e),
                        })
                n_passed  = sum(1 for r in results if r["passed"])
                pass_rate = n_passed / len(results) if results else 0.0
                valid_risks = [r["risk_score"] for r in results if r["risk_score"] is not None]
                avg_risk  = sum(valid_risks) / len(valid_risks) if valid_risks else 0.0
                store.update_test_run(run_id, results, avg_risk, pass_rate)

            async def _run_both() -> None:
                await _asyncio.gather(
                    _run_one_pipeline(pa, run_a["id"]),
                    _run_one_pipeline(pb, run_b["id"]),
                    return_exceptions=True,
                )
                if req.notify_email:
                    try:
                        ra = store.get_test_run(run_a["id"], user["id"])
                        rb = store.get_test_run(run_b["id"], user["id"])
                        from .email_sender import send_email
                        ucfg = store.get_user_config(user["id"])
                        send_email(
                            to=req.notify_email,
                            subject=f"A/B test set comparison complete — {ts['name']}",
                            body_html=(
                                f"<p><b>Pipeline A</b> ({pa['name']}): "
                                f"pass rate {(ra.get('pass_rate',0)*100):.1f}%, "
                                f"avg risk {(ra.get('avg_risk_score',0)*100):.1f}%</p>"
                                f"<p><b>Pipeline B</b> ({pb['name']}): "
                                f"pass rate {(rb.get('pass_rate',0)*100):.1f}%, "
                                f"avg risk {(rb.get('avg_risk_score',0)*100):.1f}%</p>"
                            ),
                            user_configs=ucfg,
                        )
                    except Exception:
                        pass

            _asyncio.create_task(_run_both())
            return {
                "run_a_id": run_a["id"],
                "run_b_id": run_b["id"],
                "status":   "running",
                "message":  (
                    f"A/B test set run started. "
                    f"Poll GET /user/test-runs/{run_a['id']} and /user/test-runs/{run_b['id']} for results."
                ),
            }

        # ── Test sets & regression runner ─────────────────────────────────────

        @app.get("/user/test-sets", summary="List test sets")
        async def list_test_sets(user: dict = _Depends(_jwt_user)):
            return store.get_test_sets(user["id"])

        @app.post("/user/test-sets", summary="Save a test set")
        async def save_test_set(req: TestSetSaveRequest, user: dict = _Depends(_jwt_user)):
            return store.save_test_set(user["id"], req.name, req.cases, req.pipeline_id, req.id)

        @app.get("/user/test-sets/{test_set_id}", summary="Get a test set")
        async def get_test_set(test_set_id: int, user: dict = _Depends(_jwt_user)):
            ts = store.get_test_set(test_set_id, user["id"])
            if not ts:
                raise HTTPException(404, "Test set not found.")
            return ts

        @app.delete("/user/test-sets/{test_set_id}", summary="Delete a test set")
        async def delete_test_set(test_set_id: int, user: dict = _Depends(_jwt_user)):
            ok = store.delete_test_set(test_set_id, user["id"])
            if not ok:
                raise HTTPException(404, "Test set not found.")
            return {"deleted": True}

        @app.post("/user/test-sets/{test_set_id}/run", summary="Run a regression test (async background)")
        async def run_test_set(
            test_set_id: int,
            req:  TestRunRequest,
            user: dict = _Depends(_jwt_user),
        ):
            """
            Enqueues a test run and returns immediately with run_id.
            Poll GET /user/test-runs/{run_id} for status and results.
            """
            ts = store.get_test_set(test_set_id, user["id"])
            if not ts:
                raise HTTPException(404, "Test set not found.")
            p = store.get_pipeline(req.pipeline_id, user["id"])
            if not p:
                raise HTTPException(404, "Pipeline not found.")

            run_rec      = store.create_test_run(test_set_id, user["id"], req.pipeline_id)
            run_id       = run_rec["id"]
            user_configs = store.get_user_config(user["id"])

            from .pipeline_executor import run_pipeline as _run_pipeline

            async def _bg_run():
                """Execute all test cases in background; update run record on finish."""
                cases        = ts.get("cases", [])
                results:list = []
                for case in cases:
                    query   = case.get("query", "")
                    expected= case.get("expected_output", "")
                    thresh  = float(case.get("risk_threshold", 0.65))
                    try:
                        r = await _run_pipeline(p, query, user["id"], user_configs, store)
                        risk   = r.get("risk_score", 0.0) or 0.0
                        output = r.get("result", "")
                        passed = risk < thresh
                        results.append({
                            "query": query, "expected": expected, "output": output,
                            "risk_score": risk, "risk_threshold": thresh,
                            "passed": passed, "status": r.get("status", "done"),
                        })
                    except Exception as e:
                        results.append({
                            "query": query, "expected": expected, "output": "",
                            "risk_score": None, "passed": False,
                            "status": "error", "error": str(e),
                        })

                n_passed  = sum(1 for r in results if r["passed"])
                pass_rate = n_passed / len(results) if results else 0.0
                valid_risks = [r["risk_score"] for r in results if r["risk_score"] is not None]
                avg_risk  = sum(valid_risks) / len(valid_risks) if valid_risks else 0.0
                store.update_test_run(run_id, results, avg_risk, pass_rate)

                if req.notify_email:
                    try:
                        from .email_sender import test_run_email
                        test_run_email(
                            to=req.notify_email,
                            test_set_name=ts["name"],
                            pipeline_name=p["name"],
                            pass_rate=pass_rate,
                            avg_risk=avg_risk,
                            n_cases=len(results),
                            run_id=run_id,
                            user_configs=user_configs,
                        )
                    except Exception:
                        pass

            _asyncio.create_task(_bg_run())
            return {"run_id": run_id, "status": "running",
                    "message": f"Test run started. Poll GET /user/test-runs/{run_id} for results."}

        @app.get("/user/test-sets/{test_set_id}/runs", summary="Test run history")
        async def get_test_runs(test_set_id: int, user: dict = _Depends(_jwt_user)):
            ts = store.get_test_set(test_set_id, user["id"])
            if not ts:
                raise HTTPException(404, "Test set not found.")
            return store.get_test_runs(test_set_id, user["id"])

        @app.get("/user/test-runs/{run_id}", summary="Get a single test run with results")
        async def get_test_run(run_id: int, user: dict = _Depends(_jwt_user)):
            r = store.get_test_run(run_id, user["id"])
            if not r:
                raise HTTPException(404, "Test run not found.")
            return r

        # ── Pipeline run detail (SSE resume) ──────────────────────────────────

        @app.get("/user/pipelines/{pipeline_id}/runs/{run_id}",
                 summary="Get a single pipeline run with full trace (SSE resume)")
        async def get_pipeline_run(
            pipeline_id: int,
            run_id:      int,
            user:        dict = _Depends(_jwt_user),
        ):
            """
            Returns the stored result for a completed run. Use this to resume a
            disconnected SSE stream — re-fetch the final state without re-running.
            """
            run = store.get_run_by_id(run_id, user["id"])
            if not run or run.get("pipeline_id") != pipeline_id:
                raise HTTPException(404, "Pipeline run not found.")
            return run

        # ── Automation / workflow execution log ───────────────────────────────

        @app.get("/user/automation-logs", summary="Guardrail and workflow execution audit log")
        async def get_automation_logs(
            domain: str = "",
            limit:  int = 100,
            user:   dict = _Depends(_jwt_user),
        ):
            return store.get_automation_logs(user["id"], limit=min(limit, 500), domain=domain)

        # ── Pre-built evaluators ──────────────────────────────────────────────

        @app.post("/evaluate", summary="Run pre-built evaluators on text")
        async def evaluate(req: EvaluateRequest, user: dict = _Depends(_jwt_user)):
            from .evaluators import run_evaluator, run_all, EVALUATOR_REGISTRY
            if req.evaluators:
                results = {}
                for name in req.evaluators:
                    if name not in EVALUATOR_REGISTRY:
                        raise HTTPException(400, f"Unknown evaluator '{name}'.")
                    results[name] = run_evaluator(name, req.text, req.context).to_dict()
            else:
                results = {k: v.to_dict() for k, v in run_all(req.text, req.context).items()}
            return {"results": results, "text_length": len(req.text)}

        @app.get("/evaluate/list", summary="List available evaluators")
        async def list_evaluators():
            from .evaluators import EVALUATOR_REGISTRY
            return {"evaluators": list(EVALUATOR_REGISTRY.keys())}

        # ── Org / team management ─────────────────────────────────────────────

        @app.get("/user/orgs", summary="List orgs for current user")
        async def list_orgs(user: dict = _Depends(_jwt_user)):
            return store.get_orgs_for_user(user["id"])

        @app.post("/user/orgs", summary="Create an org")
        async def create_org(req: OrgCreateRequest, user: dict = _Depends(_jwt_user)):
            return store.create_org(user["id"], req.name)

        @app.get("/user/orgs/{org_id}/members", summary="List org members")
        async def list_org_members(org_id: int, user: dict = _Depends(_jwt_user)):
            if not store.is_org_member(org_id, user["id"]):
                raise HTTPException(403, "Not a member of this org.")
            return store.get_org_members(org_id)

        @app.post("/user/orgs/{org_id}/invite", summary="Invite a user to org by email")
        async def invite_member(org_id: int, req: OrgInviteRequest, user: dict = _Depends(_jwt_user)):
            if not store.is_org_member(org_id, user["id"]):
                raise HTTPException(403, "Not a member of this org.")
            invite = store.create_invite(org_id, req.email)
            # Send invite email if SMTP configured
            user_configs = store.get_user_config(user["id"])
            if req.email and (user_configs.get("smtp_user") or os.environ.get("SMTP_USER")):
                from .email_sender import send_email
                org = store.get_orgs_for_user(user["id"])
                org_name = next((o["name"] for o in org if o["id"] == org_id), "your team")
                send_email(
                    to=req.email,
                    subject=f"You've been invited to join {org_name} on llm-guard-kit",
                    body_html=f"""<p>You've been invited to join <strong>{org_name}</strong>.</p>
<p>Click <a href="https://app.llm-guard-kit.com/app/team/accept?token={invite['token']}">here</a> to accept.</p>
<p>Token: <code>{invite['token']}</code></p>""",
                    user_configs=user_configs,
                )
            return {"invited": True, "email": req.email, "token": invite["token"],
                    "expires_at": invite["expires_at"]}

        @app.post("/user/orgs/accept", summary="Accept an org invite")
        async def accept_invite(req: OrgAcceptRequest, user: dict = _Depends(_jwt_user)):
            result = store.accept_invite(req.token, user["id"])
            if not result:
                raise HTTPException(400, "Invalid or expired invite token.")
            return {"joined": True, "org_id": result["org_id"]}

        @app.delete("/user/orgs/{org_id}/members/{member_id}", summary="Remove a member from org")
        async def remove_member(org_id: int, member_id: int, user: dict = _Depends(_jwt_user)):
            if not store.is_org_member(org_id, user["id"]):
                raise HTTPException(403, "Not a member of this org.")
            ok = store.remove_member(org_id, member_id, user["id"])
            if not ok:
                raise HTTPException(400, "Cannot remove this member (owner, or not found).")
            return {"removed": True}

        # ── LLM Proxy ─────────────────────────────────────────────────────────
        # Intercepts OpenAI/Anthropic calls, auto-scores, logs, and optionally blocks.
        # Configure SDK: openai.base_url = "http://localhost:8000/proxy/openai/v1"

        import httpx as _httpx

        async def _proxy_score_and_log(
            user_id: int, provider: str, model: str,
            pt: int, ct: int, latency_ms: int,
            question: str, output_text: str,
        ) -> None:
            """Background task: score output + log proxy request. Never raises."""
            risk_score: float | None = None
            try:
                from .label_free_scorer import LabelFreeScorer
                steps  = [{"thought": output_text, "action_type": "Finish",
                           "action_arg": output_text, "observation": ""}]
                scorer = LabelFreeScorer()
                res    = scorer.score(question, steps, output_text, finished=True)
                risk_score = res.risk_score
            except Exception:
                pass
            try:
                store.log_proxy_request(user_id, provider, model, pt, ct, risk_score, latency_ms)
            except Exception:
                pass

        @app.post("/proxy/openai/v1/chat/completions",
                  summary="OpenAI proxy — intercept, score, and forward")
        async def proxy_openai(
            request_body: dict,
            authorization: Optional[str] = Header(None),
            user: dict = _Depends(_jwt_user),
        ):
            user_configs = store.get_user_config(user["id"])
            api_key      = user_configs.get("openai_key", "")
            if not api_key:
                raise HTTPException(400, "OpenAI API key not configured in Settings.")

            t0 = time.time()
            try:
                async with _httpx.AsyncClient(timeout=60.0) as client:
                    resp = await client.post(
                        "https://api.openai.com/v1/chat/completions",
                        json=request_body,
                        headers={
                            "Authorization": f"Bearer {api_key}",
                            "Content-Type": "application/json",
                        },
                    )
                    resp.raise_for_status()
                    data = resp.json()
            except Exception as e:
                raise HTTPException(502, f"OpenAI API error: {e}")

            latency_ms  = int((time.time() - t0) * 1000)
            usage       = data.get("usage", {})
            pt          = usage.get("prompt_tokens", 0)
            ct          = usage.get("completion_tokens", 0)
            model       = data.get("model", request_body.get("model", ""))
            choices     = data.get("choices", [])
            output_text = choices[0]["message"]["content"] if choices else ""
            messages    = request_body.get("messages", [])
            question    = next((m["content"] for m in messages if m.get("role") == "user"), "")

            # Fire scoring + logging in background — don't add latency to response
            _asyncio.create_task(_proxy_score_and_log(
                user["id"], "openai", model, pt, ct, latency_ms, question, output_text
            ))
            return data

        @app.post("/proxy/anthropic/v1/messages",
                  summary="Anthropic proxy — intercept, score, and forward")
        async def proxy_anthropic(
            request_body: dict,
            authorization: Optional[str] = Header(None),
            user: dict = _Depends(_jwt_user),
        ):
            user_configs = store.get_user_config(user["id"])
            api_key      = user_configs.get("anthropic_key", "")
            if not api_key:
                raise HTTPException(400, "Anthropic API key not configured in Settings.")

            t0 = time.time()
            try:
                async with _httpx.AsyncClient(timeout=60.0) as client:
                    resp = await client.post(
                        "https://api.anthropic.com/v1/messages",
                        json=request_body,
                        headers={
                            "x-api-key": api_key,
                            "anthropic-version": "2023-06-01",
                            "Content-Type": "application/json",
                        },
                    )
                    resp.raise_for_status()
                    data = resp.json()
            except Exception as e:
                raise HTTPException(502, f"Anthropic API error: {e}")

            latency_ms  = int((time.time() - t0) * 1000)
            usage       = data.get("usage", {})
            pt          = usage.get("input_tokens", 0)
            ct          = usage.get("output_tokens", 0)
            model       = data.get("model", request_body.get("model", ""))
            content     = data.get("content", [])
            output_text = content[0].get("text", "") if content else ""
            messages    = request_body.get("messages", [])
            question    = next((m["content"] for m in messages if m.get("role") == "user"), "")

            # Fire scoring + logging in background — zero added latency
            _asyncio.create_task(_proxy_score_and_log(
                user["id"], "anthropic", model, pt, ct, latency_ms, question, output_text
            ))
            return data

        @app.get("/proxy/stats", summary="LLM proxy request stats")
        async def proxy_stats(days: int = 30, user: dict = _Depends(_jwt_user)):
            return store.get_proxy_stats(user["id"], days=min(days, 365))

        @app.get("/proxy/requests", summary="Recent proxied requests")
        async def proxy_requests(limit: int = 50, user: dict = _Depends(_jwt_user)):
            return store.get_proxy_requests(user["id"], limit=min(limit, 200))

        # ── Usage metering ────────────────────────────────────────────────────

        @app.get("/user/usage", summary="Plan usage for current billing period")
        async def get_usage(user: dict = _Depends(_jwt_user)):
            """
            Returns chains scored this month, plan limits, and percentage used.
            Used by the sidebar usage bar and onboarding flow.
            """
            from .auth import FREE_MONTHLY_LIMIT
            used  = store.get_monthly_usage(user["id"])
            plan  = user.get("plan", "free")
            limit = FREE_MONTHLY_LIMIT if plan == "free" else None
            return {
                "plan":      plan,
                "used":      used,
                "limit":     limit,
                "pct":       round(used / limit * 100, 1) if limit else None,
                "remaining": max(0, limit - used) if limit else None,
            }

        # ── Onboarding ────────────────────────────────────────────────────────

        @app.get("/user/onboarding", summary="Get onboarding progress")
        async def get_onboarding(user: dict = _Depends(_jwt_user)):
            cfg = store.get_user_config(user["id"])
            return {
                "done":         bool(cfg.get("onboarding_done")),
                "current_step": int(cfg.get("onboarding_step", 0)),
            }

        @app.post("/user/onboarding", summary="Update onboarding progress")
        async def update_onboarding(body: dict, user: dict = _Depends(_jwt_user)):
            """
            body: { step?: number, done?: bool }
            Merges into user_configs so LLM keys etc. are preserved.
            """
            update: dict = {}
            if "step" in body:
                update["onboarding_step"] = int(body["step"])
            if body.get("done"):
                update["onboarding_done"] = True
                update["onboarding_step"] = 4
            if update:
                store.set_user_config(user["id"], update)
            cfg = store.get_user_config(user["id"])
            return {
                "done":         bool(cfg.get("onboarding_done")),
                "current_step": int(cfg.get("onboarding_step", 0)),
            }

        # ── Demo scoring (onboarding wizard) ─────────────────────────────────

        @app.post("/demo/score", summary="Demo scoring — no API key required (rate-limited)")
        async def demo_score(req: ScoreRequest, user: dict = _Depends(_jwt_user)):
            """
            Score a chain for the onboarding wizard. Requires JWT (logged-in user)
            but no customer API key. Counts against monthly usage. Max 10 calls/session.
            """
            from .label_free_scorer import LabelFreeScorer
            steps  = [s.model_dump() for s in req.steps]
            result = LabelFreeScorer().score(req.question, steps, "", True)
            return {
                "risk_score":       result.risk_score,
                "needs_review":     result.needs_review,
                "behavioral_score": result.behavioral_score,
                "gmm_score":        result.gmm_score,
                "chain_id":         -1,
            }

        @app.post("/api/{domain}/score", summary="Authenticated scoring (SaaS)")
        async def saas_score(
            domain:        str,
            req:           ScoreRequest,
            authorization: Optional[str] = Header(None),
        ):
            key_info      = _auth_customer(authorization)
            scoped_domain = f"{key_info['customer_id']}:{domain}"
            steps         = [s.model_dump() for s in req.steps]

            # Plan enforcement for user-linked keys
            user = None
            user_id = key_info.get("user_id")
            if user_id is not None:
                user = store.get_user_by_id(user_id)
            if user is not None:
                from .auth import check_plan_limits
                check_plan_limits(user, store, want_judge=False)

            from .label_free_scorer import LabelFreeScorer
            scorer = LabelFreeScorer()
            pool   = store.get_calibration_pool(scoped_domain, n=200)
            if pool:
                scorer.calibrate(pool)

            result = scorer.score(req.question, steps, "", True)
            chain_id = store.add_chain(
                scoped_domain,
                {"question": req.question, "steps": steps,
                 "final_answer": "", "finished": True},
                risk_score=result.risk_score,
                alert=result.needs_review,
            )

            # Track monthly usage for user-linked keys
            if user is not None:
                store.increment_usage(user["id"])

            # Run guardrails + workflows for user-linked keys
            automation: dict = {}
            if user is not None:
                from .workflow_executor import run_automations
                automation = await run_automations(
                    user_id=user["id"],
                    domain=domain,
                    risk_score=result.risk_score,
                    question=req.question,
                    chain_id=chain_id,
                    store=store,
                )
                # Block action: return 403 immediately
                if automation.get("guardrail_action") == "block":
                    raise HTTPException(
                        status_code=403,
                        detail={
                            "error": "blocked_by_guardrail",
                            "message": "This chain was blocked by a guardrail rule.",
                            "risk_score": result.risk_score,
                        },
                    )

            return {
                "risk_score":        result.risk_score,
                "needs_review":      result.needs_review,
                "behavioral_score":  result.behavioral_score,
                "gmm_score":         result.gmm_score,
                "calibration_size":  result.calibration_size,
                "guardrail_action":  automation.get("guardrail_action"),
                "fallback_required": automation.get("fallback_required", False),
            }

        @app.get("/api/{domain}/audit-log", summary="Export audit log (SaaS)")
        async def saas_audit(
            domain:        str,
            fmt:           str = "csv",
            start:         float = 0.0,
            end:           float = 0.0,
            authorization: Optional[str] = Header(None),
        ):
            key_info      = _auth_customer(authorization)
            scoped_domain = f"{key_info['customer_id']}:{domain}"
            end           = end or time.time()
            content       = store.export_audit(scoped_domain, start, end, fmt=fmt)
            from fastapi.responses import PlainTextResponse
            media = "text/csv" if fmt == "csv" else "application/json"
            return PlainTextResponse(content, media_type=media)

    # ── Federated GMM calibration (exp68) ─────────────────────────────────────
    # Share GMM parameters (6 KB) instead of raw chains (40 KB) for
    # privacy-preserving cross-site calibration.

    @app.get("/api/{domain}/gmm-params",
             summary="Export GMM parameters for federated calibration (exp68)")
    async def export_gmm_params(
        domain:        str,
        authorization: Optional[str] = Header(None),
    ):
        """
        Export this domain's calibrated GMM parameters as JSON (~6 KB).

        The payload can be averaged across sites and re-imported via
        POST /api/{domain}/import-gmm for privacy-preserving federated
        calibration without sharing raw chain data.

        Requires: domain must be calibrated (≥20 chains scored).
        Auth: Bearer token (same as scoring endpoints).
        """
        if store is not None:
            _auth_customer(authorization)  # type: ignore[name-defined]
        scorer = _domain_scorers.get(domain)
        if scorer is None:
            # Try to warm up from store
            if store is not None:
                pool = store.get_calibration_pool(domain, n=200)
                if pool:
                    scorer = _LFS()
                    scorer.calibrate(pool)
                    _domain_scorers[domain] = scorer
        if scorer is None or scorer._gmm2 is None:
            raise HTTPException(
                404,
                f"Domain '{domain}' not calibrated. Score ≥20 chains first."
            )
        params = scorer.get_gmm_params()
        return {"domain": domain, "params": params}

    @app.post("/api/{domain}/import-gmm",
              summary="Import federated-averaged GMM parameters (exp68)")
    async def import_gmm_params(
        domain:        str,
        req:           ImportGmmRequest,
        authorization: Optional[str] = Header(None),
    ):
        """
        Import GMM parameters (possibly averaged across multiple sites) and
        apply them to this domain's scorer.  Enables calibration without
        sharing raw agent chain data.

        Exp68: federated averaging over 4 sites → AUROC within 0.02 of
        centralized training, while transferring only 762 floats (6 KB)
        vs ~40 KB of raw chains.
        """
        if store is not None:
            _auth_customer(authorization)  # type: ignore[name-defined]
        if domain not in _domain_scorers:
            _domain_scorers[domain] = _LFS()
        ok = _domain_scorers[domain].set_gmm_params(req.params)
        if not ok:
            raise HTTPException(422, "Invalid GMM params structure.")
        if store is not None:
            store.log_calibration(domain, req.n_chains or 0, auroc_estimate=0.0)
        return {
            "domain":   domain,
            "imported": True,
            "n_chains": req.n_chains,
            "message":  "GMM parameters applied. Scorer will use federated calibration.",
        }

    @app.post("/api/gmm/average",
              summary="Server-side federated averaging of GMM parameters")
    async def average_gmm_params(
        sites: list[dict],
        x_admin_key: Optional[str] = Header(None),
    ):
        """
        Weighted average of GMM parameters from multiple sites.

        Body: list of {"params": {...}, "n_chains": int} — one per site.
        Returns averaged params to push back to each site via import-gmm.

        Example use:
          1. Each site: GET /api/{domain}/gmm-params → send to central server
          2. Central: POST /api/gmm/average with all site params
          3. Each site: POST /api/{domain}/import-gmm with averaged params
        """
        if admin_key and x_admin_key != admin_key:
            raise HTTPException(401, "X-Admin-Key required.")
        if not sites:
            raise HTTPException(400, "Provide at least one site.")
        import numpy as np
        weights  = [float(s.get("n_chains", 1)) for s in sites]
        total_w  = sum(weights) or 1.0
        p0       = sites[0]["params"]
        avg: dict = {}
        for component in ("gmm2", "gmm4"):
            avg[component] = {}
            for key in ("means", "covariances", "weights"):
                arrays = [np.array(s["params"][component][key]) for s in sites]
                avg[component][key] = (
                    sum(w / total_w * a for w, a in zip(weights, arrays))
                ).tolist()
        avg["embed_model"] = p0.get("embed_model", "all-MiniLM-L6-v2")
        avg["pca_dims"]    = p0.get("pca_dims", 32)   # PCA dims (architecture constant)
        avg["n_chains"]    = int(total_w)
        return {"averaged_params": avg, "n_sites": len(sites), "total_chains": int(total_w)}

    # ── Prometheus /metrics ────────────────────────────────────────────────────

    @app.get("/metrics",
             summary="Prometheus text-format metrics for all domains")
    async def prometheus_metrics():
        """
        Expose per-domain monitoring metrics in Prometheus text format.

        Metrics exported:
          qppg_chains_total{domain}        — total chains scored
          qppg_alert_total{domain}         — total alerts fired
          qppg_alert_rate{domain}          — fraction of chains that triggered alert
          qppg_avg_risk_score{domain}      — rolling average risk score
          qppg_calibration_size{domain}    — number of calibration chains
          qppg_calibrated{domain}          — 1 if domain is calibrated, 0 otherwise

        Scrape with: prometheus.yml → scrape_configs: targets: ["localhost:8000"]
        """
        from fastapi.responses import PlainTextResponse
        lines = [
            "# HELP qppg_chains_total Total chains scored per domain",
            "# TYPE qppg_chains_total counter",
            "# HELP qppg_alert_total Total alert-level chains per domain",
            "# TYPE qppg_alert_total counter",
            "# HELP qppg_alert_rate Alert fraction (0-1) per domain",
            "# TYPE qppg_alert_rate gauge",
            "# HELP qppg_avg_risk_score Rolling average risk score per domain",
            "# TYPE qppg_avg_risk_score gauge",
            "# HELP qppg_calibration_size Number of calibration chains per domain",
            "# TYPE qppg_calibration_size gauge",
            "# HELP qppg_calibrated 1 if domain is calibrated (>=20 chains)",
            "# TYPE qppg_calibrated gauge",
            "",
        ]
        domains = []
        if store is not None:
            try:
                domains = store.get_domains()
            except Exception:
                pass
        # Include default domain from QPPGService
        default_domain = getattr(service, "domain_name", "default")
        if default_domain not in domains:
            domains.append(default_domain)

        for d in domains:
            # From SQLite store (most accurate)
            if store is not None:
                try:
                    recent = store.get_recent_risk(d, days=3650)  # all time
                    n_total = len(recent)
                    # count alerts from store
                    chains_raw = store.get_chains(d, limit=10000)
                    n_alerts = sum(1 for c in chains_raw if c.get("alert_triggered"))
                    alert_rate = n_alerts / n_total if n_total else 0.0
                    import numpy as _np
                    avg_risk = float(_np.mean([r[1] for r in recent])) if recent else 0.0
                    cal_size = len(store.get_calibration_pool(d, n=10000))
                    calibrated = 1 if cal_size >= 20 else 0
                except Exception:
                    n_total = n_alerts = 0
                    alert_rate = avg_risk = 0.0
                    cal_size = 0
                    calibrated = 0
            else:
                # Fall back to in-memory QPPGService stats for default domain
                try:
                    st = service.status()
                    n_total   = st.get("n_scored", 0)
                    n_alerts  = 0
                    alert_rate = 0.0
                    avg_risk  = 0.0
                    cal_size  = st.get("n_calibration", 0)
                    calibrated = 1 if st.get("lifecycle_stage") == "DEPLOYED" else 0
                except Exception:
                    n_total = n_alerts = 0
                    alert_rate = avg_risk = 0.0
                    cal_size = 0
                    calibrated = 0

            tag = f'{{domain="{d}"}}'
            lines += [
                f'qppg_chains_total{tag} {n_total}',
                f'qppg_alert_total{tag} {n_alerts}',
                f'qppg_alert_rate{tag} {round(alert_rate, 6)}',
                f'qppg_avg_risk_score{tag} {round(avg_risk, 6)}',
                f'qppg_calibration_size{tag} {cal_size}',
                f'qppg_calibrated{tag} {calibrated}',
                "",
            ]

        return PlainTextResponse("\n".join(lines), media_type="text/plain")

    # ── Dashboard ──────────────────────────────────────────────────────────────

    if enable_dashboard and store is not None:
        from .dashboard import add_dashboard_routes
        add_dashboard_routes(app, store, dashboard_key=dashboard_key)
    elif enable_dashboard and store is None:
        import warnings
        warnings.warn("Dashboard requires --db to be set. Dashboard disabled.")

    return app


# ── CLI entry point ───────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="llm-guard-kit HTTP API server")
    parser.add_argument("--domain",    default="default")
    parser.add_argument("--save-dir",  default=None)
    parser.add_argument("--port",      type=int, default=8000)
    parser.add_argument("--host",      default="127.0.0.1")
    parser.add_argument("--threshold", type=float, default=0.45)
    parser.add_argument("--qpd",       type=int, default=50)
    parser.add_argument("--db",        default=None,
                        help="SQLite path for SaaS/dashboard")
    parser.add_argument("--admin-key", default=None,
                        help="Secret key for /api/keys creation")
    parser.add_argument("--dashboard-key", default=None,
                        help="Secret key for /dashboard/api/admin/* routes (X-Dashboard-Key header)")
    parser.add_argument("--dashboard", action="store_true",
                        help="Enable /dashboard UI (requires --db)")
    args = parser.parse_args()

    svc = QPPGService(
        domain_name=args.domain,
        save_dir=args.save_dir,
        review_threshold=args.threshold,
        queries_per_day=args.qpd,
    )

    db_path = args.db
    if args.dashboard and not db_path:
        db_path = str(Path.home() / ".qppg" / "chains.db")

    app = create_app(
        svc,
        db_path=db_path,
        admin_key=args.admin_key,
        enable_dashboard=args.dashboard,
        dashboard_key=args.dashboard_key,
    )

    print(svc.progress_bar())
    print(f"\nStarting llm-guard-kit API on http://{args.host}:{args.port}")
    print(f"  Docs      : http://{args.host}:{args.port}/docs")
    if args.dashboard:
        print(f"  Dashboard : http://{args.host}:{args.port}/dashboard")
    print()

    import uvicorn
    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()
