"""FastAPI backend for the Grok Voice Agent web UI.

Serves the static frontend and exposes a single WebSocket endpoint (`/realtime`)
that acts as a transparent proxy to the xAI Voice Agent realtime API. The xAI
API key never leaves the server; the browser only ever talks to this app.

Run locally:
    python server.py
    # or: uvicorn server:app --reload --port 7860

Deploy: this is the entry point for a Hugging Face Space (SDK: docker).
Set the `XAI_API_KEY` secret in the Space settings.
"""

from __future__ import annotations

import asyncio
import os
from pathlib import Path

import websockets
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

load_dotenv()

# The local .env uses `xai_api_key`; HF Space secrets use `XAI_API_KEY`.
XAI_API_KEY = os.getenv("XAI_API_KEY") or os.getenv("xai_api_key")
AGENT_ID = os.getenv("XAI_AGENT_ID", "agent_otGoZoOqjIvC6jMv")
WS_URL = f"wss://api.x.ai/v1/realtime?agent_id={AGENT_ID}"

STATIC_DIR = Path(__file__).parent / "static"

app = FastAPI(title="Grok Voice Agent")


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


@app.get("/health")
async def health() -> JSONResponse:
    return JSONResponse(
        {"ok": True, "agent_id": AGENT_ID, "api_key_configured": bool(XAI_API_KEY)}
    )


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


async def _connect_upstream():
    """Open the WebSocket to xAI, tolerating the websockets header kwarg rename."""
    headers = {"Authorization": f"Bearer {XAI_API_KEY}"}
    try:
        return await websockets.connect(WS_URL, additional_headers=headers, max_size=None)
    except TypeError:
        # websockets < 12 used `extra_headers`.
        return await websockets.connect(WS_URL, extra_headers=headers, max_size=None)


@app.websocket("/realtime")
async def realtime(client: WebSocket) -> None:
    """Relay JSON events between the browser and the xAI Voice Agent."""
    await client.accept()

    if not XAI_API_KEY:
        await client.send_json(
            {"type": "error", "error": {"message": "XAI_API_KEY is not set on the server."}}
        )
        await client.close()
        return

    try:
        upstream = await _connect_upstream()
    except Exception as exc:  # noqa: BLE001 - surface any handshake failure to the UI
        await client.send_json(
            {"type": "error", "error": {"message": f"Could not reach xAI: {exc}"}}
        )
        await client.close()
        return

    async def browser_to_xai() -> None:
        try:
            while True:
                message = await client.receive_text()
                await upstream.send(message)
        except (WebSocketDisconnect, RuntimeError):
            pass

    async def xai_to_browser() -> None:
        try:
            async for message in upstream:
                await client.send_text(message)
        except websockets.exceptions.ConnectionClosed:
            pass

    up_task = asyncio.create_task(browser_to_xai())
    down_task = asyncio.create_task(xai_to_browser())
    _, pending = await asyncio.wait(
        {up_task, down_task}, return_when=asyncio.FIRST_COMPLETED
    )
    for task in pending:
        task.cancel()

    await upstream.close()
    try:
        await client.close()
    except RuntimeError:
        pass


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))