#!/usr/bin/env python3
"""Governance-Bench Agent Reasoning Audit Space.

A FastAPI application that serves as a 'try it' interface for the
governance-bench benchmark. Users can:
1. Browse tasks by category and difficulty
2. View task instructions
3. Run verify.py scripts to see expected results
4. Analyze governance reasoning with Base120 mental models
5. View leaderboard results

Designed for deployment as a Hugging Face Space.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

# Paths — use BENCH_ROOT env var if set (for Docker/HF Space), otherwise relative to file
BENCH_ROOT = Path(os.environ.get("BENCH_ROOT", str(Path(__file__).resolve().parent.parent)))
TASKS_ROOT = BENCH_ROOT / "tasks"
SCRIPTS_DIR = BENCH_ROOT / "scripts"
TEMPLATES_DIR = Path(__file__).resolve().parent / "templates"
STATIC_DIR = Path(__file__).resolve().parent / "static"

app = FastAPI(
    title="Governance-Bench Agent Reasoning Audit",
    description="Interactive benchmark for AI agent governance primitives",
    version="1.0.0",
)

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------

@dataclass
class TaskInfo:
    task_id: str
    category: str
    difficulty: str
    scoring: str
    description: str
    instruction: str
    path: Path

    @property
    def name(self) -> str:
        return self.task_id


# ---------------------------------------------------------------------------
# Task discovery
# ---------------------------------------------------------------------------

def discover_all_tasks() -> list[TaskInfo]:
    """Discover all tasks in the benchmark."""
    tasks = []
    for cat_dir in sorted(TASKS_ROOT.iterdir()):
        if not cat_dir.is_dir():
            continue
        for task_dir in sorted(cat_dir.iterdir()):
            toml = task_dir / "task.toml"
            instr = task_dir / "instruction.md"
            if not toml.exists():
                continue

            content = toml.read_text()
            difficulty = "unknown"
            scoring = "deterministic"
            description = ""

            for line in content.splitlines():
                line = line.strip()
                if line.startswith("difficulty"):
                    difficulty = line.split("=")[1].strip().strip('"')
                elif line.startswith("scoring_model"):
                    scoring = line.split("=")[1].strip().strip('"')
                elif line.startswith("description"):
                    description = line.split("=", 1)[1].strip().strip('"')

            instruction = instr.read_text() if instr.exists() else ""

            tasks.append(TaskInfo(
                task_id=task_dir.name,
                category=cat_dir.name,
                difficulty=difficulty,
                scoring=scoring,
                description=description,
                instruction=instruction,
                path=task_dir,
            ))
    return tasks


def get_task(category: str, task_id: str) -> TaskInfo | None:
    """Get a single task by category and ID."""
    task_dir = TASKS_ROOT / category / task_id
    toml = task_dir / "task.toml"
    if not toml.exists():
        return None

    content = toml.read_text()
    difficulty = "unknown"
    scoring = "deterministic"
    description = ""

    for line in content.splitlines():
        line = line.strip()
        if line.startswith("difficulty"):
            difficulty = line.split("=")[1].strip().strip('"')
        elif line.startswith("scoring_model"):
            scoring = line.split("=")[1].strip().strip('"')
        elif line.startswith("description"):
            description = line.split("=", 1)[1].strip().strip('"')

    instr = task_dir / "instruction.md"
    instruction = instr.read_text() if instr.exists() else ""

    return TaskInfo(
        task_id=task_id,
        category=category,
        difficulty=difficulty,
        scoring=scoring,
        description=description,
        instruction=instruction,
        path=task_dir,
    )


# ---------------------------------------------------------------------------
# Base120 governance analysis
# ---------------------------------------------------------------------------

BASE120_GOV_MODELS = {
    "interruptibility": {
        "name": "Interruptibility",
        "description": "Can the system be stopped at any time?",
        "primitives": ["KillSwitch", "CircuitBreaker"],
        "key_questions": [
            "What happens when EMERGENCY is engaged?",
            "Are critical tasks still allowed?",
            "Is state persisted across restarts?",
            "Can an adversarial agent bypass the switch?",
        ],
    },
    "failure_containment": {
        "name": "Failure Containment",
        "description": "Are failures isolated to prevent cascading?",
        "primitives": ["CircuitBreaker"],
        "key_questions": [
            "Does one adapter's failure cascade to others?",
            "Is there a Signal 3 detection mechanism?",
            "Can the breaker be reset without authorization?",
        ],
    },
    "scope_narrowing": {
        "name": "Monotonic Scope Narrowing",
        "description": "Do capabilities narrow as delegation chains deepen?",
        "primitives": ["DelegationTokenManager"],
        "key_questions": [
            "Can a child token widen its scope beyond the parent?",
            "Are depth bounds enforced for high-risk operations?",
            "Are delegation cycles detected and prevented?",
        ],
    },
    "authority_validation": {
        "name": "Authority Validation",
        "description": "Are all actions traceable to authorized agents?",
        "primitives": ["AuditLog", "AgentRegistry"],
        "key_questions": [
            "Are DCT entries validated against capability tokens?",
            "Are unregistered agents rejected?",
            "Can lateral authority attempts be detected?",
        ],
    },
    "data_provenance": {
        "name": "Data Provenance Tracking",
        "description": "Is the origin of every piece of data tracked?",
        "primitives": ["TaintEngine"],
        "key_questions": [
            "Can taint be laundered through intermediate tools?",
            "Are sink checks enforced for all dangerous operations?",
            "Does prompt injection propagate through the taint chain?",
        ],
    },
    "fail_closed": {
        "name": "Fail-Closed Execution",
        "description": "Does the system deny by default when uncertain?",
        "primitives": ["CapabilityFence", "OutputValidator"],
        "key_questions": [
            "What happens when a capability check fails?",
            "Are destructive commands blocked regardless of context?",
            "Is injection detection applied to all outputs?",
        ],
    },
    "drift_detection": {
        "name": "Behavioral Drift Detection",
        "description": "Can the system detect when agents drift from expected behavior?",
        "primitives": ["BehaviorMonitor", "ConvergenceDetector"],
        "key_questions": [
            "Is reward gaming detected?",
            "Are power-seeking patterns flagged?",
            "Can legitimate adaptation be distinguished from drift?",
        ],
    },
}


def analyze_task_with_base120(task: TaskInfo) -> dict[str, Any]:
    """Analyze a task using Base120 governance mental models."""
    category_map = {
        "kill_switch": "interruptibility",
        "circuit_breaker": "failure_containment",
        "delegation_chains": "scope_narrowing",
        "authority_boundaries": "authority_validation",
        "taint_tracking": "data_provenance",
        "execution_boundary": "fail_closed",
        "behavioral_drift": "drift_detection",
    }

    model_key = category_map.get(task.category, "interruptibility")
    model = BASE120_GOV_MODELS.get(model_key, {})

    return {
        "task_id": task.task_id,
        "category": task.category,
        "base120_model": model.get("name", "Unknown"),
        "model_key": model_key,
        "description": model.get("description", ""),
        "primitives": model.get("primitives", []),
        "key_questions": model.get("key_questions", []),
        "difficulty": task.difficulty,
        "scoring": task.scoring,
        "regulatory_mapping": {
            "EU_AI_Act": "Article 14 (Human Oversight)" if task.category in ["kill_switch", "circuit_breaker"] else "Article 12 (Logging)",
            "NIST_AI_RMF": "MANAGE" if task.category in ["kill_switch", "circuit_breaker", "execution_boundary"] else "MEASURE",
        },
    }


# ---------------------------------------------------------------------------
# Task execution
# ---------------------------------------------------------------------------

def run_task_verification(task: TaskInfo, timeout: int = 30) -> dict[str, Any]:
    """Run a task's verify.py and return results."""
    verify_path = task.path / "tests" / "verify.py"

    if not verify_path.exists():
        return {"error": "verify.py not found", "exit_code": -1}

    start = time.time()
    try:
        result = subprocess.run(
            [sys.executable, str(verify_path)],
            cwd=str(task.path),
            timeout=timeout,
            capture_output=True,
            text=True,
        )
        duration = round(time.time() - start, 2)

        # Read reward and dimension scores
        reward = 0.0
        reward_file = task.path / "reward.txt"
        if reward_file.exists():
            try:
                reward = float(reward_file.read_text().strip())
            except ValueError:
                pass

        dims = {}
        dim_file = task.path / "dimension_scores.json"
        if dim_file.exists():
            try:
                dims = json.loads(dim_file.read_text())
            except json.JSONDecodeError:
                pass

        return {
            "exit_code": result.returncode,
            "overall_score": reward,
            "dimensions": dims,
            "duration_sec": duration,
            "stdout": result.stdout[-2000:] if result.stdout else "",
            "stderr": result.stderr[-500:] if result.stderr else "",
        }
    except subprocess.TimeoutExpired:
        return {"error": "Timeout", "exit_code": -1, "duration_sec": timeout}
    except Exception as e:
        return {"error": str(e), "exit_code": -1}


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    """Main page — benchmark overview."""
    tasks = discover_all_tasks()
    categories = {}
    for t in tasks:
        cat = categories.setdefault(t.category, {
            "name": t.category,
            "tasks": [],
            "easy": 0, "medium": 0, "hard": 0,
            "deterministic": 0, "rubric": 0,
        })
        cat["tasks"].append(t)
        cat[t.difficulty] = cat.get(t.difficulty, 0) + 1
        cat[t.scoring] = cat.get(t.scoring, 0) + 1

    return templates.TemplateResponse(
        request,
        "index.html",
        {
            "categories": categories,
            "total_tasks": len(tasks),
            "base120_models": BASE120_GOV_MODELS,
        },
    )


@app.get("/api/tasks", response_class=JSONResponse)
async def api_tasks(category: str | None = None):
    """List all tasks, optionally filtered by category."""
    tasks = discover_all_tasks()
    if category:
        tasks = [t for t in tasks if t.category == category]
    return [
        {
            "task_id": t.task_id,
            "category": t.category,
            "difficulty": t.difficulty,
            "scoring": t.scoring,
            "description": t.description,
        }
        for t in tasks
    ]


@app.get("/api/task/{category}/{task_id}", response_class=JSONResponse)
async def api_task_detail(category: str, task_id: str):
    """Get task details including instruction and Base120 analysis."""
    task = get_task(category, task_id)
    if not task:
        return JSONResponse({"error": "Task not found"}, status_code=404)

    analysis = analyze_task_with_base120(task)
    return {
        "task_id": task.task_id,
        "category": task.category,
        "difficulty": task.difficulty,
        "scoring": task.scoring,
        "description": task.description,
        "instruction": task.instruction,
        "base120_analysis": analysis,
    }


@app.post("/api/task/{category}/{task_id}/run", response_class=JSONResponse)
async def api_run_task(category: str, task_id: str):
    """Run a task's verification script and return results."""
    task = get_task(category, task_id)
    if not task:
        return JSONResponse({"error": "Task not found"}, status_code=404)

    result = run_task_verification(task)
    analysis = analyze_task_with_base120(task)

    return {
        "task_id": task.task_id,
        "category": task.category,
        "result": result,
        "base120_analysis": analysis,
    }


@app.get("/api/base120/models", response_class=JSONResponse)
async def api_base120_models():
    """List all Base120 governance mental models."""
    return BASE120_GOV_MODELS


@app.get("/api/leaderboard", response_class=JSONResponse)
async def api_leaderboard():
    """Get leaderboard data from results files."""
    results = []
    results_dir = BENCH_ROOT / "results"
    if results_dir.exists():
        for f in sorted(results_dir.glob("*.json")):
            try:
                data = json.loads(f.read_text())
                results.append({
                    "agent": data.get("agent", "unknown"),
                    "model": data.get("model", "unknown"),
                    "timestamp": data.get("timestamp", ""),
                    "overall_score": data.get("overall_score", 0.0),
                    "category_scores": data.get("category_scores", {}),
                })
            except json.JSONDecodeError:
                continue

    return sorted(results, key=lambda x: x["overall_score"], reverse=True)


@app.get("/health")
async def health():
    return {"status": "healthy", "tasks": len(discover_all_tasks())}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)