"""
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 sys
import time
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, Header, HTTPException
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 = ""


# ── 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,
) -> 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
    """
    app = FastAPI(
        title="llm-guard-kit API",
        description=(
            "Real-time confidence scoring and monitoring for LLM agent chains. "
            "v0.3.0 — SaaS auth, per-domain SQLite, monitoring dashboard."
        ),
        version="0.3.0",
    )

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

    # ── 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.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}

    # ── 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):
            with store._conn() as conn:
                rows = conn.execute(
                    "SELECT * FROM api_keys WHERE key_hash LIKE ?",
                    (key_id + "%",),
                ).fetchall()
            if not rows:
                raise HTTPException(404, "Key not found.")
            row = dict(rows[0])
            row.pop("key_hash", None)
            return row

        @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]

            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)
            store.add_chain(
                scoped_domain,
                {"question": req.question, "steps": steps,
                 "final_answer": "", "finished": True},
                risk_score=result.risk_score,
                alert=result.needs_review,
            )
            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,
            }

        @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)

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

    if enable_dashboard and store is not None:
        from .dashboard import add_dashboard_routes
        add_dashboard_routes(app, store)
    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", 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,
    )

    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()
