app/main.py

"""FastAPI application: realtime security-log RAG, CSV ingest, mock emitter.

Run (self-host):  uvicorn app.main:app --host 0.0.0.0 --port 8000
Realtime needs a single worker so all WebSocket clients share one process.
"""
from __future__ import annotations

from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

from fastapi import FastAPI, File, Query, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from . import __version__
from .config import get_settings
from .db import fetch_alerts, fetch_events, fetch_stats, init_db
from .generator import generate, generate_live, to_csv
from .ingest import ingest_rows, parse_csv
from .rag import ask as rag_ask
from .realtime import manager, stream

ROOT = Path(__file__).resolve().parents[1]
WEB_DIR = ROOT / "web"
HTML_DIR = ROOT / "html"

settings = get_settings()


@asynccontextmanager
async def lifespan(app: FastAPI):
    await init_db()
    yield
    await stream.stop()


app = FastAPI(title="ROGLAG Realtime", version=__version__, lifespan=lifespan)
app.add_middleware(
    CORSMiddleware, allow_origins=["*"], allow_methods=["*"],
    allow_headers=["*"], allow_credentials=False,
)


# ---------------------------------------------------------------- broadcasting
async def broadcast_events(events: list[dict[str, Any]], summary: dict[str, Any], kind: str) -> None:
    """Push events (capped) + a summary frame to all WebSocket clients."""
    cap = 300
    to_push = events if len(events) <= cap else (
        [e for e in events if e.get("severity") == "critical"][:cap] or events[:cap]
    )
    for ev in to_push:
        await manager.broadcast({"type": "event", "event": _slim(ev)})
    await manager.broadcast({"type": "summary", "kind": kind, "summary": summary})


def _slim(ev: dict[str, Any]) -> dict[str, Any]:
    return {
        "ts": ev.get("timestamp") or ev.get("ts"),
        "source": ev.get("source"),
        "category": ev.get("category"),
        "label": ev.get("label"),
        "severity": ev.get("severity"),
        "mitre": ev.get("mitre"),
        "risk_score": ev.get("risk_score"),
        "src_ip": ev.get("src_ip"),
        "dst_ip": ev.get("dst_ip"),
        "user": ev.get("user") or ev.get("username"),
        "message": ev.get("message"),
        "reasons": ev.get("reasons"),
        "expected_category": ev.get("expected_category"),
        "matches_label": ev.get("matches_label"),
    }


async def emit_mock(count: int, profile: str) -> dict[str, Any]:
    rows = generate_live(count, profile=profile)
    events, summary = await ingest_rows(rows)
    await broadcast_events(events, summary, kind="mock")
    return summary


# ---------------------------------------------------------------- API: health/stats
@app.get("/health")
async def health() -> dict[str, Any]:
    return {
        "ok": True, "version": __version__,
        "storage": "postgres" if settings.is_postgres else "sqlite",
        "llm_ready": settings.llm_ready,
        "ws_clients": manager.count,
        "stream_running": stream.is_running(),
    }


@app.get("/api/stats")
async def api_stats() -> Any:
    return await fetch_stats()


@app.get("/api/logs")
async def api_logs(limit: int = 100, offset: int = 0, category: str = "",
                   severity: str = "", source: str = "", q: str = "") -> Any:
    return await fetch_events(limit, offset, category, severity, source, q)


@app.get("/api/alerts")
async def api_alerts(limit: int = 20) -> Any:
    return await fetch_alerts(limit)


class AskBody(BaseModel):
    question: str
    limit: int = 8
    use_llm: bool = False
    model: str = ""


@app.post("/api/ask")
async def api_ask(body: AskBody) -> Any:
    if not body.question.strip():
        return JSONResponse({"error": "question is required"}, status_code=400)
    return await rag_ask(body.question.strip(), min(body.limit, 20), body.use_llm, body.model)


# ---------------------------------------------------------------- CSV upload (realtime detect)
@app.post("/api/ingest/csv")
async def api_ingest_csv(file: UploadFile = File(...)) -> Any:
    data = await file.read()
    rows = parse_csv(data)
    if not rows:
        return JSONResponse({"error": "no parseable rows (need a 'message' or 'source' column)"}, status_code=400)
    if len(rows) > settings.max_csv_rows:
        rows = rows[: settings.max_csv_rows]
    events, summary = await ingest_rows(rows)
    summary["filename"] = file.filename
    await broadcast_events(events, summary, kind="csv")
    return summary


# ---------------------------------------------------------------- CSV generator (unique each call)
@app.get("/api/generate/csv")
async def api_generate_csv(count: int = Query(200, ge=1, le=20000),
                           profile: str = Query("mixed")) -> Response:
    records = generate(count, profile=profile, seed=None)  # unseeded -> unique each time
    csv_text = to_csv(records)
    fname = f"roglag_logs_{profile}_{len(records)}.csv"
    return Response(
        content=csv_text, media_type="text/csv",
        headers={"Content-Disposition": f'attachment; filename="{fname}"'},
    )


# ---------------------------------------------------------------- mock emitter (realtime)
@app.post("/api/mock/emit")
async def api_mock_emit(count: int = Query(20, ge=1), profile: str = Query("mixed")) -> Any:
    count = min(count, settings.max_mock_burst)
    return await emit_mock(count, profile)


@app.post("/api/mock/stream/start")
async def api_stream_start(rate: int = Query(3, ge=1, le=50),
                           interval: float = Query(1.0, ge=0.2, le=10.0),
                           profile: str = Query("mixed")) -> Any:
    await stream.start(emit_mock, rate=min(rate, settings.max_mock_burst), interval=interval, profile=profile)
    return {"running": stream.is_running(), "rate": rate, "interval": interval, "profile": profile}


@app.post("/api/mock/stream/stop")
async def api_stream_stop() -> Any:
    await stream.stop()
    return {"running": stream.is_running()}


# ---------------------------------------------------------------- WebSocket
@app.websocket("/ws")
async def ws_feed(ws: WebSocket) -> None:
    await manager.connect(ws)
    try:
        await ws.send_json({"type": "hello", "clients": manager.count})
        while True:
            await ws.receive_text()  # keep-alive; we ignore client messages
    except WebSocketDisconnect:
        await manager.disconnect(ws)
    except Exception:
        await manager.disconnect(ws)


# ---------------------------------------------------------------- static / pages
@app.get("/")
async def index() -> Response:
    live = WEB_DIR / "live.html"
    return FileResponse(live) if live.exists() else JSONResponse({"ok": True, "see": "/docs-html"})


@app.get("/docs-html")
async def docs_html() -> Response:
    return FileResponse(HTML_DIR / "docs.html")


@app.get("/plan")
async def plan_html() -> Response:
    return FileResponse(HTML_DIR / "plan.html")


if WEB_DIR.exists():
    app.mount("/web", StaticFiles(directory=str(WEB_DIR)), name="web")
if HTML_DIR.exists():
    app.mount("/html", StaticFiles(directory=str(HTML_DIR)), name="html")