"""FastAPI application — serves UI and WebSocket stream."""
from __future__ import annotations

import asyncio
import json
import logging
import os
import threading
from pathlib import Path

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, HTMLResponse, Response
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

from agentradar.server import events as ev_module
from agentradar.server import websocket as ws_module
from agentradar.storage import snapshots as snap_module
from agentradar.storage import failures as fail_module
from agentradar.storage import config as cfg_module

logger = logging.getLogger("agentradar.server")

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

app = FastAPI(title="agentradar", version="1.0.0")

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

_NO_CACHE = {"Cache-Control": "no-store, no-cache, must-revalidate", "Pragma": "no-cache"}


@app.get("/", response_class=HTMLResponse)
async def root():
    html_path = FRONTEND_DIR / "index.html"
    if html_path.exists():
        return HTMLResponse(html_path.read_text(encoding="utf-8"), headers=_NO_CACHE)
    return HTMLResponse("<h1>agentradar frontend not found</h1>", status_code=500)


@app.get("/js/{filename:path}")
async def serve_js(filename: str):
    path = FRONTEND_DIR / "js" / filename
    if not path.exists():
        return Response("Not found", status_code=404)
    return FileResponse(str(path), media_type="application/javascript", headers=_NO_CACHE)


@app.get("/css/{filename:path}")
async def serve_css(filename: str):
    path = FRONTEND_DIR / "css" / filename
    if not path.exists():
        return Response("Not found", status_code=404)
    return FileResponse(str(path), media_type="text/css", headers=_NO_CACHE)


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws_module.connect(ws)
    # Flush all pending queue events into _run_state BEFORE syncing to new client.
    # This fixes the race where the pipeline finishes before the browser connects.
    await _flush_queue_once()
    # Send full run state to the newly-connected client
    state = _run_state.copy()
    if state:
        try:
            await ws.send_text(json.dumps({"event_type": "state_sync", "data": state}, default=str))
        except Exception:
            pass
    try:
        while True:
            data = await ws.receive_text()
            if data == "ping":
                await ws.send_text("pong")
    except WebSocketDisconnect:
        pass
    finally:
        await ws_module.disconnect(ws)


async def _flush_queue_once():
    """Drain every event currently sitting in the queue (non-blocking)."""
    import queue as q_mod
    q = ev_module._event_queue
    loop = asyncio.get_event_loop()
    while True:
        try:
            event = q.get_nowait()
        except q_mod.Empty:
            break
        try:
            update_run_state(event)
            await ws_module.broadcast(event)
        except Exception:
            pass


# In-memory run state for new connections (state sync)
_run_state: dict = {}


def update_run_state(event: dict) -> None:
    """Update in-memory state for late-joining clients."""
    etype = event.get("event_type", "")
    if etype == "run_start":
        _run_state.clear()
        _run_state["run_id"] = event.get("run_id")
        _run_state["run_name"] = event.get("run_id", "")
        _run_state["status"] = "live"
        _run_state["agents"] = {}
        _run_state["edges"] = []
        _run_state["steps"] = []
        _run_state["frameworks"] = event.get("frameworks", [])
        _run_state["total_tokens"] = 0
        _run_state["total_cost"] = 0.0
    elif etype in ("agent_start", "agent_end", "agent_error"):
        agents = _run_state.setdefault("agents", {})
        agent_id = event.get("agent_id", "")
        existing = agents.get(agent_id, {})
        agents[agent_id] = {
            **existing,
            **event,
            # Preserve input captured at agent_start — later events send input=None
            "input": existing.get("input") if event.get("input") is None else event.get("input"),
        }
        if etype == "agent_end":
            agents[agent_id]["status"] = "success"
        elif etype == "agent_error":
            agents[agent_id]["status"] = "failed"
        elif etype == "agent_start":
            agents[agent_id]["status"] = "running"
        # Track edges (parent_id can be str or list[str] for merge points)
        parent_id = event.get("parent_id")
        if parent_id and parent_id != agent_id:
            parent_ids = parent_id if isinstance(parent_id, list) else [parent_id]
            for pid in parent_ids:
                edge = {"source": pid, "target": agent_id}
                if edge not in _run_state.get("edges", []):
                    _run_state.setdefault("edges", []).append(edge)
        # Track steps
        _run_state.setdefault("steps", []).append(event)
        # Accumulate tokens
        tokens = event.get("tokens", {})
        _run_state["total_tokens"] = _run_state.get("total_tokens", 0) + tokens.get("total", 0)
    elif etype == "run_end":
        _run_state["status"] = "complete"


# Background task: drain queue and broadcast
async def _drain_queue():
    loop = asyncio.get_event_loop()
    q = ev_module._event_queue
    while True:
        try:
            # Non-blocking get from thread-safe queue
            event = await loop.run_in_executor(None, _get_nowait, q)
            if event is None:
                await asyncio.sleep(0.05)
                continue
            update_run_state(event)
            await ws_module.broadcast(event)
            # Save failure immediately on agent_error (real-time), and also on run_end
            if event.get("event_type") in ("agent_error", "run_end"):
                _check_and_save_failure()
        except Exception as e:
            logger.debug("Queue drain error: %s", e)
            await asyncio.sleep(0.1)


def _get_nowait(q: "asyncio.Queue") -> dict | None:
    import queue as q_mod
    try:
        return q.get_nowait()
    except q_mod.Empty:
        return None


def _check_and_save_failure():
    agents = _run_state.get("agents", {})
    has_failure = any(a.get("status") == "failed" for a in agents.values())
    if has_failure:
        try:
            fail_module.save_failure(_run_state.copy())
        except Exception as e:
            logger.debug("Failed to save failure: %s", e)


@app.on_event("startup")
async def startup():
    asyncio.create_task(_drain_queue())
    logger.info("agentradar server started")


# REST endpoints for the UI
@app.get("/api/runs/current")
async def get_current_run():
    return _run_state


@app.get("/api/failures")
async def get_failures():
    return fail_module.list_failures()


@app.get("/api/failures/{failure_id}")
async def get_failure(failure_id: str):
    return fail_module.load_failure(failure_id)


@app.delete("/api/failures")
async def clear_failures():
    fail_module.clear_failures()
    return {"ok": True}


@app.get("/api/snapshots")
async def list_snapshots():
    return snap_module.list_snapshots()


@app.post("/api/snapshots")
async def save_snapshot(name: str = "snapshot"):
    path = snap_module.save_snapshot(_run_state.copy(), name)
    return {"path": str(path), "ok": True}


@app.get("/api/snapshots/{name}")
async def load_snapshot(name: str):
    return snap_module.load_snapshot(name)


@app.get("/api/config")
async def get_config():
    return cfg_module.load_config()


@app.post("/api/config")
async def save_config(config: dict):
    cfg_module.save_config(config)
    return {"ok": True}
