"""FastAPI Backend Server with SSE Real-Time Telemetry and Model Orchestration API."""
import asyncio
import json
import logging
from typing import Dict, Any, Optional, List
from fastapi import FastAPI, Request, HTTPException, Query
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

import ram_optimizer
import vram_arbitrator

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("model_manager_server")

app = FastAPI(
    title="HyperSwap // GPU Program Swapper & Telemetry API",
    version="1.0.0",
    description="High-performance VRAM arbitration and 64GB RAM cache orchestrator for simultaneous Ollama and ComfyUI workloads on Linux.",
    docs_url="/docs",
    redoc_url="/redoc",
)

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

# Pydantic Request Models
class SwitchRequest(BaseModel):
    model: str = Field(..., description="Name of the Ollama model to hot-swap to in VRAM", example="qwen3.8fast:latest")
    keep_alive: Optional[str] = Field("30m", description="Keep-alive duration in VRAM (e.g. 5m, 30m, 0)", example="30m")

class WarmRequest(BaseModel):
    model_name: Optional[str] = Field(None, description="Ollama model name to warm into OS page cache", example="gemma4:26b")
    filepath: Optional[str] = Field(None, description="Absolute file path of Safetensors/GGUF to warm into RAM", example="/home/drjones/ComfyUI/models/checkpoints/v1-5-pruned-emaonly-fp16.safetensors")

class BenchmarkRequest(BaseModel):
    iterations: Optional[int] = Field(2, description="Number of back-and-forth switch iterations to measure", example=2)
    models: Optional[List[str]] = Field(None, description="Optional pair of models to benchmark between", example=["qwen3.8fast:latest", "smtek/Qwen3.8-27B:Q2_K_XL"])


# ==========================================
# REST API ENDPOINTS
# ==========================================

@app.get("/api/stats", summary="Full System Snapshot", tags=["Telemetry"])
async def get_all_stats() -> Dict[str, Any]:
    """Gather complete live snapshot of GPU hardware, host RAM, Ollama, ComfyUI, and switch history."""
    gpu_stats = vram_arbitrator.get_gpu_hardware_stats()
    mem_stats = ram_optimizer.get_detailed_meminfo()
    ollama_state = await vram_arbitrator.get_ollama_live_state()
    comfy_state = await vram_arbitrator.get_comfyui_live_state()
    history = vram_arbitrator.get_switch_history()
    comfy_models = ram_optimizer.find_comfy_model_files()
    
    return {
        "timestamp": asyncio.get_event_loop().time(),
        "gpu": gpu_stats,
        "ram": mem_stats,
        "ollama": ollama_state,
        "comfyui": comfy_state,
        "history": history,
        "comfy_models_count": len(comfy_models),
    }

@app.get("/api/gpu", summary="GPU Sensors and VRAM Breakdown", tags=["Telemetry"])
async def get_gpu_metrics() -> Dict[str, Any]:
    """Retrieve detailed NVML sensors (utilization %, temp, power, fan, clocks, and per-process VRAM allocation)."""
    return vram_arbitrator.get_gpu_hardware_stats()

@app.get("/api/memory", summary="Host RAM and Page Cache Breakdown", tags=["Telemetry"])
async def get_ram_metrics() -> Dict[str, Any]:
    """Retrieve precise host 64GB DDR5 RAM breakdown, active cache size, and cache hit ratios."""
    return ram_optimizer.get_detailed_meminfo()

@app.get("/api/stream", summary="Real-Time SSE Telemetry Stream", tags=["Telemetry"])
async def sse_telemetry_stream(request: Request):
    """Server-Sent Events (SSE) streaming real-time statistics at 1Hz for dynamic dashboards."""
    async def event_generator():
        while True:
            if await request.is_disconnected():
                break
            try:
                stats = await get_all_stats()
                yield f"data: {json.dumps(stats)}\n\n"
            except Exception as e:
                logger.error(f"SSE stream error: {e}")
                yield f"data: {json.dumps({'error': str(e)})}\n\n"
            await asyncio.sleep(1.0)
            
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        }
    )

@app.post("/api/switch-model", summary="Hot-Swap Ollama LLM in VRAM", tags=["Orchestration"])
async def api_switch_model(req: SwitchRequest):
    """Hot-swap the active Ollama model in VRAM and measure exact load duration and token evaluation speed."""
    res = await vram_arbitrator.switch_ollama_model(req.model, keep_alive=req.keep_alive or "30m")
    if not res.get("success"):
        raise HTTPException(status_code=500, detail=res.get("error"))
    return res

@app.post("/api/free-vram", summary="Soft-Yield Ollama VRAM", tags=["Orchestration"])
async def api_free_vram():
    """Instruct Ollama to instantly yield VRAM to 0MB in ~15ms while preserving model weights in the 64GB host RAM page cache."""
    return await vram_arbitrator.instant_free_ollama_vram()

@app.post("/api/comfy-free", summary="Purge ComfyUI VRAM Cache", tags=["Orchestration"])
async def api_comfy_free():
    """Purge loaded diffusion models and VRAM cache from the ComfyUI pipeline."""
    return await vram_arbitrator.instant_free_comfyui_vram()

@app.post("/api/warm-all", summary="Pre-warm All Models into RAM Cache", tags=["Memory Optimization"])
async def api_warm_all():
    """Pre-fault and read all installed Ollama GGUF models and ComfyUI Safetensors checkpoints into the Linux OS Page Cache."""
    return await ram_optimizer.warm_all_models()

@app.post("/api/warm-model", summary="Pre-warm Single Model or File", tags=["Memory Optimization"])
async def api_warm_model(req: WarmRequest):
    """Pre-warm a specific Ollama model or individual file path into Linux RAM cache."""
    if req.model_name:
        return await ram_optimizer.warm_ollama_model(req.model_name, keep_alive="1m")
    elif req.filepath:
        return ram_optimizer.warm_file_to_ram(req.filepath)
    else:
        raise HTTPException(status_code=400, detail="model_name or filepath required")

@app.get("/api/models", summary="List All Installed Models", tags=["Catalog"])
async def api_get_models():
    """List all installed Ollama models and discovered ComfyUI model checkpoints/safetensors on disk with sizes and quantization levels."""
    ollama_state = await vram_arbitrator.get_ollama_live_state()
    comfy_models = ram_optimizer.find_comfy_model_files()
    return {
        "ollama_models": ollama_state.get("installed_models", []),
        "comfy_models": comfy_models,
    }

@app.get("/api/history", summary="Model Switch History Log", tags=["Analytics"])
async def api_get_history(limit: int = Query(20, description="Max history items to return")):
    """Get the recent history log of model switch events, swap durations (in ms), and RAM cache hit status."""
    history = vram_arbitrator.get_switch_history()
    return history[:limit]

@app.post("/api/benchmark", summary="Run Latency Benchmark", tags=["Analytics"])
async def api_run_benchmark(req: BenchmarkRequest):
    """Run an automated benchmark swapping between available models to measure round-trip latency and RAM cache effectiveness."""
    from mcp_server import run_model_switch_benchmark
    res_str = await run_model_switch_benchmark(iterations=req.iterations or 2)
    return json.loads(res_str)

# Mount static web UI files
app.mount("/static", StaticFiles(directory="/home/drjones/unified-model-manager/static"), name="static")

@app.get("/", summary="Dashboard Web UI", tags=["UI"])
async def root_index():
    with open("/home/drjones/unified-model-manager/static/index.html", "r") as f:
        content = f.read()
    return HTMLResponse(content=content)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=9090, reload=False, log_level="info")