from pathlib import Path

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from app.config import security_config
from app.handlers.secure import handle_secure
from app.handlers.vulnerable import handle_vulnerable
from app.models.scenarios import load_attack_samples, resolve_attack_meta
from app.models.schemas import (
    AttackMeta,
    ChatRequest,
    ChatResult,
    CompareResponse,
    SecurityConfigUpdate,
)
from app.security.layer import security_layer

FRONTEND_DIR = Path(__file__).resolve().parent.parent / "frontend"

app = FastAPI(title="LLM 安全攻防演示平台", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

if FRONTEND_DIR.exists():
    app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static")


def _resolve_task(scenario_id: str, task: str | None) -> str:
    if task:
        return task
    scenario = load_attack_samples().get(scenario_id, {})
    return scenario.get("task", "faq")


@app.get("/")
async def index():
    index_file = FRONTEND_DIR / "index.html"
    if index_file.exists():
        return FileResponse(index_file)
    return {"message": "LLM Security Demo API", "docs": "/docs"}


@app.get("/api/scenarios")
async def get_scenarios():
    return load_attack_samples()


@app.get("/api/security/config")
async def get_security_config():
    return security_config.model_dump()


@app.put("/api/security/config")
async def update_security_config(body: SecurityConfigUpdate):
    data = body.model_dump(exclude_none=True)
    for key, value in data.items():
        setattr(security_config, key, value)
    return security_config.model_dump()


@app.get("/api/audit/logs")
async def get_audit_logs():
    return security_layer.audit_logs[-50:]


@app.post("/api/chat", response_model=ChatResult)
async def chat(req: ChatRequest):
    task = _resolve_task(req.scenario_id, req.task)
    if req.mode == "vulnerable":
        return await handle_vulnerable(req.message, scenario_id=req.scenario_id, task=task)
    return await handle_secure(req.message, scenario_id=req.scenario_id, task=task)


@app.post("/api/compare", response_model=CompareResponse)
async def compare(req: ChatRequest):
    task = _resolve_task(req.scenario_id, req.task)
    meta = resolve_attack_meta(req.scenario_id, req.sample_id)

    vulnerable = await handle_vulnerable(req.message, scenario_id=req.scenario_id, task=task)
    secure = await handle_secure(req.message, scenario_id=req.scenario_id, task=task)

    return CompareResponse(
        attack_meta=AttackMeta(**meta),
        vulnerable=vulnerable,
        secure=secure,
    )


@app.post("/api/redteam/run")
async def redteam_run():
    samples_data = load_attack_samples()
    report = []

    for scenario_id, scenario in samples_data.items():
        task = scenario.get("task", "faq")
        for sample in scenario.get("samples", []):
            result = await compare(
                ChatRequest(
                    message=sample["input"],
                    scenario_id=scenario_id,
                    sample_id=sample["id"],
                )
            )
            report.append(
                {
                    "sample_id": sample["id"],
                    "scenario": scenario_id,
                    "level": sample.get("level", "未知"),
                    "vulnerable_success": result.vulnerable.attack_success,
                    "secure_blocked": not result.secure.attack_success,
                    "blocked_by": result.secure.blocked_by,
                }
            )

    total = len(report)
    blocked = sum(1 for r in report if r["secure_blocked"])
    vuln_hit = sum(1 for r in report if r["vulnerable_success"])

    by_level: dict[str, dict] = {}
    for r in report:
        lv = r["level"]
        if lv not in by_level:
            by_level[lv] = {"total": 0, "blocked": 0}
        by_level[lv]["total"] += 1
        if r["secure_blocked"]:
            by_level[lv]["blocked"] += 1

    return {
        "summary": {
            "total": total,
            "secure_blocked": blocked,
            "vulnerable_hit": vuln_hit,
            "block_rate": f"{blocked / total * 100:.1f}%" if total else "0%",
            "by_level": {
                k: {
                    **v,
                    "block_rate": f"{v['blocked'] / v['total'] * 100:.0f}%",
                }
                for k, v in by_level.items()
            },
        },
        "results": report,
    }