src/server/app.py

"""FastAPI app: WebSocket broadcast, REST /stats, Prometheus /metrics, static dashboard."""

import asyncio
import json
import logging
import logging.config
import time
from collections import deque
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles

from src.config import settings
from src.pipeline.runner import run_pipeline
from src.stream.generator import generate_stream

# Structured JSON logging
logging.config.dictConfig({
    "version": 1,
    "disable_existing_loggers": False,
    "formatters": {
        "json": {
            "()": "logging.Formatter",
            "fmt": '{"time":"%(asctime)s","level":"%(levelname)s","logger":"%(name)s","message":"%(message)s"}',
            "datefmt": "%Y-%m-%dT%H:%M:%S",
        }
    },
    "handlers": {
        "console": {"class": "logging.StreamHandler", "formatter": "json", "stream": "ext://sys.stdout"}
    },
    "root": {"level": "INFO", "handlers": ["console"]},
})

logger = logging.getLogger(__name__)

REPLAY_BUFFER_SIZE = settings.server.replay_buffer_size

# Shared state updated by the pipeline, read by /stats and /metrics
STATE: dict = {
    "total_observations": 0,
    "total_anomalies_detected": 0,
    "total_drift_events": 0,
    "current_precision": 0.0,
    "current_recall": 0.0,
    "current_f1": 0.0,
    "cycle_start_time": 0.0,
}

replay_buffer: deque = deque(maxlen=REPLAY_BUFFER_SIZE)

_pipeline_task: asyncio.Task | None = None


class ConnectionManager:
    def __init__(self) -> None:
        self._connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket) -> None:
        await websocket.accept()
        self._connections.append(websocket)

    def register(self, websocket: WebSocket) -> None:
        self._connections.append(websocket)

    def disconnect(self, websocket: WebSocket) -> None:
        if websocket in self._connections:
            self._connections.remove(websocket)

    async def broadcast(self, message: dict) -> None:
        text = json.dumps(message)
        dead: list[WebSocket] = []
        for ws in self._connections:
            try:
                await ws.send_text(text)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)


manager = ConnectionManager()


async def _broadcast_with_replay(message: dict) -> None:
    replay_buffer.append(message)
    await manager.broadcast(message)


def _reset_state() -> None:
    STATE["total_observations"] = 0
    STATE["total_anomalies_detected"] = 0
    STATE["total_drift_events"] = 0
    STATE["current_precision"] = 0.0
    STATE["current_recall"] = 0.0
    STATE["current_f1"] = 0.0
    STATE["cycle_start_time"] = time.time()


async def _run_pipeline_loop() -> None:
    cfg_s = settings.stream
    cfg_d = settings.detector
    cfg_dr = settings.drift
    while True:
        _reset_state()
        logger.info("cycle_start", extra={"seed": cfg_s.seed})
        stream = generate_stream(
            phase_a_length=cfg_s.phase_a_length,
            phase_b_length=cfg_s.phase_b_length,
            phase_c_length=cfg_s.phase_c_length,
            drift_magnitude=cfg_s.drift_magnitude,
            anomaly_rate=cfg_s.anomaly_rate,
            point_ratio=cfg_s.point_ratio,
            delay=cfg_s.delay,
            seed=cfg_s.seed,
        )
        await run_pipeline(
            stream,
            _broadcast_with_replay,
            state=STATE,
            anomaly_threshold=cfg_d.threshold,
            n_trees=cfg_d.n_trees,
            height=cfg_d.height,
            window_size=cfg_d.window_size,
            drift_delta=cfg_dr.delta,
            drift_grace_period=cfg_dr.grace_period,
        )


def _ensure_pipeline_running() -> None:
    global _pipeline_task
    if _pipeline_task is None or _pipeline_task.done():
        _pipeline_task = asyncio.create_task(_run_pipeline_loop())


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield
    global _pipeline_task
    if _pipeline_task is not None and not _pipeline_task.done():
        _pipeline_task.cancel()
        try:
            await _pipeline_task
        except asyncio.CancelledError:
            pass


app = FastAPI(title="Real-Time Anomaly Detection", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET"],
    allow_headers=["*"],
)

STATIC_DIR = Path(__file__).resolve().parent / "static"
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/")
async def root():
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/dashboard")
async def dashboard():
    return FileResponse(STATIC_DIR / "index.html")


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> None:
    await websocket.accept()
    for msg in replay_buffer:
        try:
            await websocket.send_text(json.dumps(msg))
        except Exception:
            break
    manager.register(websocket)
    _ensure_pipeline_running()
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        manager.disconnect(websocket)


@app.get("/stats")
async def stats():
    return {
        "total_observations": STATE["total_observations"],
        "total_anomalies_detected": STATE["total_anomalies_detected"],
        "total_drift_events": STATE["total_drift_events"],
        "current_precision": STATE["current_precision"],
        "current_recall": STATE["current_recall"],
        "current_f1": STATE["current_f1"],
    }


@app.get("/metrics", response_class=PlainTextResponse)
async def metrics():
    """Prometheus-compatible metrics endpoint."""
    uptime = time.time() - STATE["cycle_start_time"] if STATE["cycle_start_time"] else 0.0
    lines = [
        "# HELP anomaly_observations_total Total observations processed in current cycle",
        "# TYPE anomaly_observations_total counter",
        f'anomaly_observations_total {STATE["total_observations"]}',
        "",
        "# HELP anomaly_alerts_total Total anomaly alerts fired in current cycle",
        "# TYPE anomaly_alerts_total counter",
        f'anomaly_alerts_total {STATE["total_anomalies_detected"]}',
        "",
        "# HELP anomaly_drift_events_total Total drift events detected in current cycle",
        "# TYPE anomaly_drift_events_total counter",
        f'anomaly_drift_events_total {STATE["total_drift_events"]}',
        "",
        "# HELP anomaly_precision Running precision (alerts that were true anomalies)",
        "# TYPE anomaly_precision gauge",
        f'anomaly_precision {STATE["current_precision"]:.6f}',
        "",
        "# HELP anomaly_recall Running recall (true anomalies that were detected)",
        "# TYPE anomaly_recall gauge",
        f'anomaly_recall {STATE["current_recall"]:.6f}',
        "",
        "# HELP anomaly_f1 Running F1 score",
        "# TYPE anomaly_f1 gauge",
        f'anomaly_f1 {STATE["current_f1"]:.6f}',
        "",
        "# HELP anomaly_cycle_uptime_seconds Seconds since current stream cycle started",
        "# TYPE anomaly_cycle_uptime_seconds gauge",
        f"anomaly_cycle_uptime_seconds {uptime:.3f}",
        "",
    ]
    return "\n".join(lines)


def main() -> None:
    import os
    import uvicorn
    port = int(os.environ.get("PORT", str(settings.server.port)))
    uvicorn.run("src.server.app:app", host="0.0.0.0", port=port, reload=False)


if __name__ == "__main__":
    main()