"""Fabric Control Plane — registry + placement + scheduler + installer."""
from __future__ import annotations
import asyncio
import logging
import time

from fastapi import FastAPI, HTTPException, Query, Body
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from fabric.shared.schemas import (
    Heartbeat, PlaceRequest, PlaceResponse,
    RegistryModelsResponse, RegistryNodesResponse,
)
from fabric.control.registry import Registry
from fabric.control.scheduler import Scheduler
from fabric.control.installer import Installer, RECIPES
from fabric.control import cuda_checkpoint

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
log = logging.getLogger("fabric.control")

app = FastAPI(title="Fabric Control Plane", version="0.2.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

registry = Registry()
scheduler = Scheduler(registry)
installer = Installer()
_start_time = time.time()

# Additive VRAM-capacity reporting: per-node physical-GPU capacity captured from
# heartbeats, keyed by node_id. Kept outside the strict registry schemas so the
# change stays backward-compatible (absent fields degrade to 0 / []).
_node_capacity: dict[str, dict] = {}


class GpuMemInfo(BaseModel):
    """Per-physical-GPU memory capacity entry (mirrors the agent payload)."""
    gpu_index: int = 0
    mem_total_mb: int = 0
    mem_used_mb: int = 0


class HeartbeatExt(Heartbeat):
    """Heartbeat + additive VRAM-capacity fields. Old agents omit these and the
    defaults apply; updated agents carry per-GPU capacity through to the registry."""
    gpus: list[GpuMemInfo] = []
    vram_total_mb: int = 0


@app.get("/health")
async def health():
    return {
        "ok": True,
        "version": registry.version,
        "nodes": len(registry.get_alive_nodes()),
        "modules": len(registry.get_all_modules()),
        "uptime_s": int(time.time() - _start_time),
        "cuda_checkpoint_available": cuda_checkpoint.is_available(),
    }

@app.post("/registry/heartbeat")
async def heartbeat(hb: HeartbeatExt):
    registry.handle_heartbeat(hb)
    _node_capacity[hb.node_id] = {
        "vram_total_mb": int(hb.vram_total_mb or 0),
        "gpus": [g.model_dump(mode="json") for g in hb.gpus],
    }
    return {"ok": True, "version": registry.version}

@app.get("/registry/models")
async def list_models():
    modules = registry.get_all_modules()
    alive = registry.get_alive_nodes()
    base = RegistryModelsResponse(
        models=modules,
        nodes_count=len(alive),
        total_gpu_vram_mb=registry.total_gpu_vram_mb(),
    )
    data = base.model_dump(mode="json")
    data["total_gpu_capacity_mb"] = sum(
        int((_node_capacity.get(n.node_id) or {}).get("vram_total_mb", 0) or 0)
        for n in alive
    )
    return data

@app.get("/registry/nodes")
async def list_nodes():
    nodes_out = []
    for n in registry.get_alive_nodes():
        d = n.model_dump(mode="json")
        cap = _node_capacity.get(n.node_id) or {}
        d["vram_total_mb"] = int(cap.get("vram_total_mb", 0) or 0)
        d["gpus"] = cap.get("gpus", [])
        nodes_out.append(d)
    return {"nodes": nodes_out}

@app.get("/registry/hypervisors")
async def list_hypervisors():
    """Group nodes by hypervisor — multi-hypervisor visibility."""
    by_hyp: dict[str, list[dict]] = {}
    for n in registry.get_alive_nodes():
        hyp = n.hypervisor or "unknown"
        by_hyp.setdefault(hyp, []).append({
            "node_id": n.node_id,
            "node_ip": n.node_ip,
            "modules": len(n.modules),
            "gpu_processes": len(n.gpu_status),
            "vram_used_mb": sum(g.vram_mb for g in n.gpu_status),
        })
    return {"hypervisors": by_hyp, "total_hypervisors": len(by_hyp)}

@app.post("/place", response_model=PlaceResponse)
async def place(req: PlaceRequest):
    result = registry.place(req)
    if result is None:
        raise HTTPException(status_code=404, detail=f"Model '{req.model}' not found in any node")
    if result.state in ("unloaded", "error"):
        await scheduler.ensure_loaded(result.node_id, result.module_name)
    return result

@app.post("/admin/evict")
async def admin_evict(node_id: str, module_name: str, reason: str = "admin"):
    ok = await scheduler.evict(node_id, module_name, reason)
    if not ok:
        raise HTTPException(400, "Eviction failed")
    return {"ok": True}

@app.post("/admin/load")
async def admin_load(node_id: str, module_name: str):
    ok = await scheduler.ensure_loaded(node_id, module_name)
    if not ok:
        raise HTTPException(400, "Load failed")
    return {"ok": True}

@app.get("/scheduler/utilization")
async def scheduler_utilization():
    return scheduler.get_utilization()

@app.post("/scheduler/preempt")
async def scheduler_preempt(model: str, reason: str = "placement"):
    result = await scheduler.preempt_for(model, reason)
    if result is None:
        raise HTTPException(503, "No resources available for preemption")
    return result

@app.post("/batch")
async def batch_submit(model: str, items: list[dict] = Body(...), priority: str = "normal"):
    return await scheduler.submit_batch(model, items, priority)

@app.get("/batch/{job_id}")
async def batch_status(job_id: str):
    status = scheduler.get_batch_status(job_id)
    if status is None:
        raise HTTPException(404, "Batch job not found")
    return status

# --- v2: Installer endpoints ---

@app.get("/install/recipes")
async def list_recipes():
    return {"recipes": installer.list_recipes()}

@app.post("/install")
async def install_model(hf_repo: str, recipe: str = "vllm", name: str | None = None):
    try:
        return await installer.install(hf_repo, recipe, name)
    except ValueError as e:
        raise HTTPException(400, str(e))

@app.get("/install/{job_id}")
async def install_status(job_id: str):
    s = installer.get_status(job_id)
    if s is None:
        raise HTTPException(404, "Install job not found")
    return s

@app.get("/install")
async def list_install_jobs():
    return {"jobs": installer.list_jobs()}

# --- v2: cuda-checkpoint endpoints ---

@app.get("/checkpoint/availability")
async def checkpoint_availability():
    return {"available": cuda_checkpoint.is_available(), "binary": cuda_checkpoint.CUDA_CHECKPOINT_BIN}

@app.post("/checkpoint/{pid}/toggle")
async def checkpoint_toggle(pid: int, timeout_s: int = 30):
    ok, msg = await cuda_checkpoint.checkpoint(pid, timeout_s)
    return {"ok": ok, "pid": pid, "message": msg}

@app.get("/checkpoint/{pid}/state")
async def checkpoint_state(pid: int):
    state = await cuda_checkpoint.get_state(pid)
    return {"pid": pid, "state": state}


async def _prune_loop():
    while True:
        await asyncio.sleep(15)
        stale = registry.prune_stale_nodes()
        if stale:
            log.warning("Pruned stale nodes: %s", stale)

async def _idle_sweep_loop():
    while True:
        await asyncio.sleep(30)
        await scheduler.idle_sweep()

@app.on_event("startup")
async def startup():
    asyncio.create_task(_prune_loop())
    asyncio.create_task(_idle_sweep_loop())
    log.info("Fabric Control Plane v0.2.0 started on port 7700")