"""
token_server.py
Serves:
  GET /token?room=<name>  → LiveKit JWT token
  GET /models             → list of Ollama/LM Studio models
  static/index.html       → browser UI
"""

import os, json, asyncio, httpx
from pathlib import Path
from dotenv import load_dotenv
from fastapi import FastAPI, Query
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
import uvicorn
from livekit.api import AccessToken, VideoGrants, RoomConfiguration, RoomAgentDispatch

load_dotenv()

LIVEKIT_API_KEY    = os.getenv("LIVEKIT_API_KEY", "devkey")
LIVEKIT_API_SECRET = os.getenv("LIVEKIT_API_SECRET", "secret")
OLLAMA_BASE_URL    = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")

app = FastAPI()

@app.get("/token")
async def get_token(
    room: str = Query("my-room"),
    identity: str = Query("browser-user"),
    model: str = Query(""),
    prompt: str = Query(""),
):
    attributes = {}
    if model:
        attributes["model"] = model
    if prompt:
        attributes["prompt"] = prompt

    token = (
        AccessToken(LIVEKIT_API_KEY, LIVEKIT_API_SECRET)
        .with_identity(identity)
        .with_name(identity)
        .with_grants(VideoGrants(room_join=True, room=room))
        .with_room_config(RoomConfiguration(
            agents=[RoomAgentDispatch(agent_name="")]
        ))
    )
    if attributes:
        token = token.with_attributes(attributes)

    return {"token": token.to_jwt(), "room": room}

@app.get("/models")
async def list_models():
    """Proxy to Ollama/LM Studio model list endpoint."""
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            # OpenAI-compat endpoint (works for Ollama & LM Studio)
            r = await client.get(OLLAMA_BASE_URL.rstrip("/v1").rstrip("/") + "/api/tags")
            if r.status_code == 200:
                data = r.json()
                models = [m["name"] for m in data.get("models", [])]
                return {"models": models}
    except Exception:
        pass
    # Fallback: OpenAI /v1/models
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            r = await client.get(OLLAMA_BASE_URL.rstrip("/") + "/models")
            if r.status_code == 200:
                data = r.json()
                models = [m["id"] for m in data.get("data", [])]
                return {"models": models}
    except Exception:
        pass
    return {"models": [], "error": "Could not reach local model server"}

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

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7731)