"""Phase B: FastAPI backend with WebSocket streaming for KernBench viewer."""

import asyncio
import json
import pathlib
import webbrowser
import threading

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
import uvicorn

STATIC_DIR = pathlib.Path(__file__).parent / "static"
DEFAULT_PORT = 8765

app = FastAPI(title="KernBench Viewer")


# ── REST: workloads + events ───────────────────────────────────────
@app.get("/api/workloads")
async def list_workloads():
    from kernbench.sim_engine.event_log import get_workload_list
    return JSONResponse(content=get_workload_list())


@app.get("/api/events/demo")
async def get_demo_events():
    from kernbench.sim_engine.event_log import generate_demo_events
    return JSONResponse(content=generate_demo_events())


@app.get("/api/events/{workload_id}")
async def get_workload_events(workload_id: str):
    from kernbench.sim_engine.event_log import generate_workload_events
    try:
        return JSONResponse(content=generate_workload_events(workload_id))
    except ValueError as e:
        return JSONResponse(content={"error": str(e)}, status_code=404)


# ── WebSocket: stream events with timing ───────────────────────────
@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            msg = await ws.receive_text()
            data = json.loads(msg)
            cmd = data.get("cmd")

            if cmd == "load_demo" or cmd == "load_workload":
                from kernbench.sim_engine.event_log import generate_workload_events
                wid = data.get("workload_id", "probe-h2d")
                events = generate_workload_events(wid)
                await ws.send_json({"type": "events", "events": events})

            elif cmd == "stream_demo" or cmd == "stream_workload":
                from kernbench.sim_engine.event_log import generate_workload_events
                wid = data.get("workload_id", "probe-h2d")
                events = generate_workload_events(wid)
                speed = data.get("speed", 1.0)  # ns → real-time multiplier

                # Send metadata first
                duration = max(e["t_ns"] for e in events) if events else 0
                await ws.send_json({
                    "type": "stream_start",
                    "total_events": len(events),
                    "duration_ns": duration,
                })

                prev_t = 0.0
                for ev in events:
                    # Wait proportional to simulation time gap
                    dt = ev["t_ns"] - prev_t
                    if dt > 0 and speed > 0:
                        # Scale: 1 ns sim → (1/speed) ms real
                        await asyncio.sleep(dt / (speed * 1000.0))
                    await ws.send_json({"type": "event", **ev})
                    prev_t = ev["t_ns"]

                await ws.send_json({"type": "stream_end"})

            elif cmd == "ping":
                await ws.send_json({"type": "pong"})

    except WebSocketDisconnect:
        pass
    except Exception:
        pass


# ── Static files (must be last) ────────────────────────────────────
@app.get("/")
async def index():
    return FileResponse(STATIC_DIR / "index.html")


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


# ── Entry point (called from CLI) ──────────────────────────────────
def serve(port: int = DEFAULT_PORT, open_browser: bool = True) -> None:
    import socket

    # Check port availability before starting
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(("127.0.0.1", port))
        except OSError:
            print(f"ERROR: Port {port} is already in use.")
            print(f"  Try: kernbench web --port {port + 1}")
            print(f"  Or kill the existing process using port {port}")
            return

    url = f"http://127.0.0.1:{port}"
    print(f"KernBench Viewer: {url}")

    if open_browser:
        threading.Timer(0.5, webbrowser.open, args=(url,)).start()

    uvicorn.run(app, host="127.0.0.1", port=port, log_level="warning")


if __name__ == "__main__":
    serve()
