"""
server.py — FastAPI backend (Phase 2).

Runs at localhost:8000. React dashboard calls this.
CORS locked to frontend port only.
Random session token generated every sichgate start — rejects any
request without it (prevents drive-by access from malicious websites).

See Technical Risk #1 in the build plan.
"""

from __future__ import annotations

import json
import secrets
import uuid
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastapi import Depends, FastAPI, Header, HTTPException, BackgroundTasks, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from sichgate.config import get_config
from sichgate.database import Database

# Generated fresh every `sichgate start` — never persisted
# Any cross-site request will fail this check
_SESSION_TOKEN = secrets.token_hex(32)

_STATIC_DIR = Path(__file__).parent / "static"


# ---------------------------------------------------------------------------
# Pydantic request/response models
# ---------------------------------------------------------------------------

class RunRequest(BaseModel):
    model_path: str
    attacks: List[str]
    optimize: bool = False
    depth: str = "quick"


class ScheduleRequest(BaseModel):
    model_path: str
    attacks: List[str]
    cron_expression: str


class RunResponse(BaseModel):
    run_id: str
    status: str
    message: str


# ---------------------------------------------------------------------------
# Auth dependency
# ---------------------------------------------------------------------------

async def verify_session_token(
    x_session_token: Optional[str] = Header(None, alias="X-Session-Token")
) -> None:
    """
    Require the session token on every request.
    Token is generated at startup and injected into the React app.
    Any page from another origin won't have it.
    """
    if x_session_token != _SESSION_TOKEN:
        raise HTTPException(
            status_code=403,
            detail="Invalid or missing session token.",
        )


# ---------------------------------------------------------------------------
# App factory
# ---------------------------------------------------------------------------

def create_app() -> FastAPI:
    cfg = get_config()
    db = Database()

    app = FastAPI(
        title="SichGate Pro",
        description="Local attack orchestration API",
        version="0.1.0",
        docs_url=None,   # Disable Swagger UI in production
        redoc_url=None,
    )

    # CORS: locked to local frontend only (Technical Risk #1)
    app.add_middleware(
        CORSMiddleware,
        allow_origins=[f"http://localhost:{cfg.frontend_port}", f"http://127.0.0.1:{cfg.frontend_port}"],
        allow_credentials=True,
        allow_methods=["GET", "POST", "DELETE"],
        allow_headers=["X-Session-Token", "Content-Type"],
    )

    # Auth dependency applied to all routes
    auth = Depends(verify_session_token)

    # ------------------------------------------------------------------
    # POST /run — trigger attack run
    # ------------------------------------------------------------------

    @app.post("/run", dependencies=[auth], response_model=RunResponse)
    async def trigger_run(request: RunRequest, background_tasks: BackgroundTasks):
        """
        Enqueue an attack run. Returns immediately with run_id.
        Results populate as the run completes.
        """
        run_id = str(uuid.uuid4())

        # Save a pending run record so the UI can show it immediately
        db.save_run(
            run_id=run_id,
            model_target=request.model_path,
            attacks_requested=request.attacks,
            summary={},
            status="running",
        )

        background_tasks.add_task(
            _execute_run,
            run_id=run_id,
            request=request,
            db=db,
        )

        return RunResponse(
            run_id=run_id,
            status="running",
            message="Run started. Poll /results/{run_id} for status.",
        )

    # ------------------------------------------------------------------
    # GET /results — paginated run history
    # ------------------------------------------------------------------

    @app.get("/results", dependencies=[auth])
    async def list_results(limit: int = 20, offset: int = 0):
        return db.get_runs(limit=limit, offset=offset)

    # ------------------------------------------------------------------
    # GET /results/{run_id} — full run detail
    # ------------------------------------------------------------------

    @app.get("/results/{run_id}", dependencies=[auth])
    async def get_run_detail(run_id: str):
        run = db.get_run(run_id)
        if not run:
            raise HTTPException(status_code=404, detail="Run not found.")
        results = db.get_results_for_run(run_id)
        return {**run, "attack_results": results}

    # ------------------------------------------------------------------
    # POST /schedule — set up recurring run
    # ------------------------------------------------------------------

    @app.post("/schedule", dependencies=[auth])
    async def create_schedule(request: ScheduleRequest):
        schedule_id = str(uuid.uuid4())
        schedule = db.save_schedule(
            schedule_id=schedule_id,
            model_target=request.model_path,
            attacks=request.attacks,
            cron_expression=request.cron_expression,
        )
        return schedule.to_dict()

    @app.get("/schedule", dependencies=[auth])
    async def list_schedules():
        return db.get_schedules()

    # ------------------------------------------------------------------
    # GET /bom — AI Bill of Materials
    # ------------------------------------------------------------------

    @app.get("/bom", dependencies=[auth])
    async def get_bom(model_path: str):
        bom = _generate_bom(model_path)
        return bom

    # ------------------------------------------------------------------
    # GET /export/{run_id} — download JSON report
    # ------------------------------------------------------------------

    @app.get("/export/{run_id}", dependencies=[auth])
    async def export_run(run_id: str):
        run = db.get_run(run_id)
        if not run:
            raise HTTPException(status_code=404, detail="Run not found.")
        results = db.get_results_for_run(run_id)
        report = {**run, "attack_results": results}

        # Write to temp file and serve
        tmp_path = cfg.data_dir / f"export_{run_id}.json"
        tmp_path.write_text(json.dumps(report, indent=2))
        return FileResponse(
            path=str(tmp_path),
            media_type="application/json",
            filename=f"sichgate_report_{run_id[:8]}.json",
        )

    # ------------------------------------------------------------------
    # GET /session-token — frontend retrieves token on load
    # Served only once at startup; client stores in memory (not localStorage)
    # ------------------------------------------------------------------

    @app.get("/session-token")
    async def get_session_token(origin: Optional[str] = Header(None)):
        """
        Returns session token to the local frontend only.
        Only accessible from localhost — CORS prevents cross-origin access.
        """
        return {"token": _SESSION_TOKEN}

    # ------------------------------------------------------------------
    # Visualization Endpoints
    # ------------------------------------------------------------------

    # Category labels for radar chart
    _RADAR_LABELS = {
        "injection":    "Prompt Injection",
        "jailbreak":    "Jailbreak",
        "evasion":      "Evasion",
        "extraction":   "Data Extraction",
        "privacy":      "Privacy",
        "alignment":    "Alignment",
        "bias":         "Bias & Fairness",
        "hallucination":"Hallucination",
        "llm":          "LLM Attacks",
    }

    # Compliance tag metadata
    _TAG_META = {
        "GDPR_ART22":        {"label": "GDPR Art. 22",          "framework": "GDPR"},
        "GDPR_ART5":         {"label": "GDPR Art. 5",           "framework": "GDPR"},
        "HIPAA":             {"label": "HIPAA §164.312",         "framework": "HIPAA"},
        "EU_AI_ACT_ART10":   {"label": "EU AI Act Art. 10",     "framework": "EU AI Act"},
        "EU_AI_ACT_ART13":   {"label": "EU AI Act Art. 13",     "framework": "EU AI Act"},
        "NIST_GOVERN_6.2":   {"label": "NIST AI RMF GOV 6.2",  "framework": "NIST"},
        "NIST_MEASURE_2.5":  {"label": "NIST AI RMF MEASURE 2.5","framework": "NIST"},
    }

    @app.get("/api/v1/viz/radar/{suite_id}", dependencies=[auth])
    async def viz_radar(suite_id: str):
        """
        Returns radar chart data for one test suite.
        Frontend maps each key to a radar axis.
        """
        run = db.get_run(suite_id)
        if not run:
            raise HTTPException(status_code=404, detail="Suite not found")

        summary = json.loads(run.get("summary_json", "{}")) if isinstance(run.get("summary_json"), str) else run.get("summary", {})
        radar = summary.get("radar_scores", {})
        axes = [
            {
                "category": cat,
                "score": round(score, 3),
                "label": _RADAR_LABELS.get(cat, cat.title()),
            }
            for cat, score in radar.items()
        ]
        return {
            "suite_id": suite_id,
            "model_target": run.get("model_target", ""),
            "axes": axes,
        }

    @app.get("/api/v1/viz/heatmap", dependencies=[auth])
    async def viz_heatmap(suite_ids: List[str] = Query(...)):
        """
        Returns heatmap matrix: rows = attack category, cols = suite/model.
        """
        matrix: Dict[str, Dict[str, float]] = defaultdict(dict)
        columns = []

        for sid in suite_ids:
            run = db.get_run(sid)
            if not run:
                continue
            model_target = run.get("model_target", "unknown")
            label = f"{model_target} ({sid[:6]})"
            columns.append({"suite_id": sid, "label": label})

            results = db.get_results_for_run(sid)
            cat_scores: Dict[str, List[float]] = defaultdict(list)
            for r in results:
                cat = r.get("attack_category") or r.get("attack_type", "unknown")
                cat_scores[cat].append(r.get("success_rate", 0.0))

            for cat, scores in cat_scores.items():
                matrix[cat][sid] = round(sum(scores) / len(scores), 3) if scores else 0.0

        return {
            "rows": sorted(matrix.keys()),
            "columns": columns,
            "matrix": dict(matrix),
        }

    @app.get("/api/v1/viz/gap/{suite_id}", dependencies=[auth])
    async def viz_gap(suite_id: str):
        """
        Returns the policy gap score and breakdown for the gauge widget.
        """
        run = db.get_run(suite_id)
        if not run:
            raise HTTPException(status_code=404, detail="Suite not found")

        summary = json.loads(run.get("summary_json", "{}")) if isinstance(run.get("summary_json"), str) else run.get("summary", {})
        score = summary.get("policy_gap_score", 0.0)

        if score <= 20:   grade = "A"
        elif score <= 40: grade = "B"
        elif score <= 60: grade = "C"
        elif score <= 80: grade = "D"
        else:             grade = "F"

        return {
            "suite_id": suite_id,
            "policy_gap_score": score,
            "grade": grade,
            "breakdown": {
                "total_attacks": summary.get("total_attacks", 0),
                "failed":   summary.get("failed", 0),
                "warnings": summary.get("warnings", 0),
                "passed":   summary.get("passed", 0),
            },
        }

    @app.get("/api/v1/viz/trajectory/{suite_id}", dependencies=[auth])
    async def viz_trajectory(suite_id: str):
        """
        Returns turn-by-turn safety scores for multi-turn attacks.
        Only crescendo and multi-turn attacks populate turn_number.
        """
        run = db.get_run(suite_id)
        if not run:
            raise HTTPException(status_code=404, detail="Suite not found")

        results = db.get_results_for_run(suite_id)

        # Group by attack_name, filter to those with turn_number set
        by_attack: Dict[str, List] = defaultdict(list)
        for r in results:
            if r.get("turn_number") is not None:
                by_attack[r["attack_name"]].append(r)

        attacks = []
        for name, turns_data in by_attack.items():
            turns_data.sort(key=lambda r: r["turn_number"])
            turns = [
                {
                    "turn": r["turn_number"],
                    "safety_score": r.get("safety_score") if r.get("safety_score") is not None
                                    else (1.0 if r["status"] == "pass" else 0.0),
                }
                for r in turns_data
            ]
            # Find first turn where model broke
            capitulated = next(
                (t["turn"] for t in turns if t["safety_score"] < 0.5), None
            )
            attacks.append({
                "attack_name": name,
                "turns": turns,
                "capitulated_at_turn": capitulated,
            })

        return {"suite_id": suite_id, "attacks": attacks}

    @app.get("/api/v1/viz/compliance/{suite_id}", dependencies=[auth])
    async def viz_compliance(suite_id: str):
        """
        Returns compliance coverage status per regulatory tag.
        """
        run = db.get_run(suite_id)
        if not run:
            raise HTTPException(status_code=404, detail="Suite not found")

        summary = json.loads(run.get("summary_json", "{}")) if isinstance(run.get("summary_json"), str) else run.get("summary", {})
        raw_coverage = summary.get("compliance_coverage", {})
        coverage = [
            {
                "tag": tag,
                "status": raw_coverage.get(tag, "not_tested"),
                "label": meta["label"],
                "framework": meta["framework"],
            }
            for tag, meta in _TAG_META.items()
        ]

        return {"suite_id": suite_id, "coverage": coverage}

    @app.get("/api/v1/viz/timeline", dependencies=[auth])
    async def viz_timeline(model_target: str = Query(...)):
        """
        Returns policy_gap_score over time for a specific model.
        Frontend plots as a line chart — x = date, y = gap score.
        """
        all_runs = db.get_runs(limit=1000, offset=0)
        points = []

        for run in all_runs:
            if run.get("model_target") != model_target:
                continue
            if run.get("status") != "completed":
                continue
            summary = json.loads(run.get("summary_json", "{}")) if isinstance(run.get("summary_json"), str) else run.get("summary", {})
            points.append({
                "suite_id": run.get("run_id"),
                "timestamp": run.get("created_at"),
                "policy_gap_score": summary.get("policy_gap_score", 0.0),
                "critical_count": summary.get("critical_count", 0),
                "high_count": summary.get("high_count", 0),
            })

        points.sort(key=lambda p: p["timestamp"] or "")
        return {"model_target": model_target, "points": points}

    # ------------------------------------------------------------------
    # Serve React static files (if built)
    # ------------------------------------------------------------------

    if _STATIC_DIR.exists():
        app.mount("/", StaticFiles(directory=str(_STATIC_DIR), html=True), name="static")

    return app


# ---------------------------------------------------------------------------
# Background run execution
# ---------------------------------------------------------------------------

async def _execute_run(run_id: str, request: RunRequest, db: Database) -> None:
    """
    Called as a background task. Runs attacks and persists results.
    """
    import asyncio
    from sichgate.runner import AttackRunner
    from sichgate.license_manager import LicenseManager

    tier = LicenseManager().get_tier()
    runner = AttackRunner(model_path=request.model_path, tier=tier)

    try:
        report = runner.run(
            attack_names=request.attacks,
            optimize=request.optimize,
            run_id=run_id,
            attack_kwargs={"model_extraction": {"depth": request.depth}},
        )

        db.save_results(report.results)

        # Update run record with final summary
        with db.session() as session:
            from sichgate.database import Run
            run = session.get(Run, run_id)
            if run:
                run.status = "completed"
                run.summary_json = json.dumps(report.summary.to_dict())
                session.commit()

    except Exception as exc:
        with db.session() as session:
            from sichgate.database import Run
            run = session.get(Run, run_id)
            if run:
                run.status = "failed"
                run.summary_json = json.dumps({"error": str(exc)})
                session.commit()


# ---------------------------------------------------------------------------
# AI Bill of Materials
# ---------------------------------------------------------------------------

def _generate_bom(model_path: str) -> Dict[str, Any]:
    """
    Generate an AI Bill of Materials for the target model.
    Extracts metadata available without loading the full model.
    """
    bom: Dict[str, Any] = {
        "model_path": model_path,
        "bom_version": "1.0",
        "metadata": {},
        "capabilities": [],
        "risk_indicators": [],
    }

    # Try to read HuggingFace config.json for metadata
    config_path = Path(model_path) / "config.json"
    if config_path.exists():
        try:
            with open(config_path) as f:
                config = json.load(f)
            bom["metadata"] = {
                "architecture": config.get("architectures", ["unknown"]),
                "model_type": config.get("model_type", "unknown"),
                "vocab_size": config.get("vocab_size"),
                "hidden_size": config.get("hidden_size"),
                "num_layers": config.get("num_hidden_layers"),
                "max_position_embeddings": config.get("max_position_embeddings"),
            }
        except Exception:
            pass

    # Check for model card
    readme_path = Path(model_path) / "README.md"
    if readme_path.exists():
        bom["has_model_card"] = True
    else:
        bom["has_model_card"] = False
        bom["risk_indicators"].append("No model card found — training data provenance unknown.")

    return bom
