"""Fabric Node Agent — module discovery, GPU monitoring, systemd lifecycle, heartbeat."""
from __future__ import annotations
import asyncio
import logging
import os
import socket
import subprocess
import time

from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import httpx

from fabric.shared.schemas import Heartbeat, ModuleState
from fabric.agent.discovery import scan_modules
from fabric.agent.gpu_monitor import get_per_pid_vram, get_gpu_count
from fabric.agent.systemd_wrap import SystemdWrapRuntime
from fabric.agent.heartbeat import heartbeat_loop

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

NODE_ID = os.environ.get("FABRIC_NODE_ID", socket.gethostname())
NODE_IP = os.environ.get("FABRIC_NODE_IP", "")
AGENT_PORT = int(os.environ.get("FABRIC_AGENT_PORT", "7701"))
HYPERVISOR = os.environ.get("FABRIC_HYPERVISOR", "unknown")


class GpuMemInfo(BaseModel):
    """Per-physical-GPU memory capacity entry (additive VRAM-capacity reporting)."""
    gpu_index: int = 0
    mem_total_mb: int = 0
    mem_used_mb: int = 0


class HeartbeatExt(Heartbeat):
    """Heartbeat + additive VRAM-capacity fields. Subclassing keeps the base
    payload identical for control planes that haven't been updated yet (they
    ignore the extra fields) while carrying per-GPU capacity to updated ones."""
    gpus: list[GpuMemInfo] = []
    vram_total_mb: int = 0


app = FastAPI(title=f"Fabric Agent ({NODE_ID})", version="0.1.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

_runtimes: dict[str, SystemdWrapRuntime] = {}
_start_time = time.time()


def _detect_ip() -> str:
    if NODE_IP:
        return NODE_IP
    try:
        r = subprocess.run(["hostname", "-I"], capture_output=True, text=True, timeout=5)
        for ip in r.stdout.strip().split():
            if not ip.startswith("127."):
                return ip
    except Exception:
        pass
    return "127.0.0.1"


def _get_gpu_mem_info() -> list[dict]:
    """Per-physical-GPU total/used VRAM via nvidia-smi (one entry per GPU, counted
    once — NOT per process). Robust if nvidia-smi is missing or fails -> returns []."""
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,memory.total,memory.used",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5,
        )
    except Exception as e:
        log.warning("nvidia-smi capacity query failed: %s", e)
        return []
    if r.returncode != 0:
        return []
    gpus: list[dict] = []
    for line in r.stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 2:
            continue
        try:
            idx = int(parts[0])
            total = int(float(parts[1]))
        except ValueError:
            continue
        used = 0
        if len(parts) >= 3:
            try:
                used = int(float(parts[2]))
            except ValueError:
                used = 0
        gpus.append({"gpu_index": idx, "mem_total_mb": total, "mem_used_mb": used})
    return gpus


def _load_modules():
    global _runtimes
    manifests = scan_modules()
    for m in manifests:
        if m.runtime.value == "systemd-wrap" and m.systemd:
            _runtimes[m.name] = SystemdWrapRuntime(m)
            log.info("Loaded systemd-wrap: %s -> %s", m.name, m.systemd.systemd_unit)
        else:
            log.warning("Unsupported runtime %s for %s (v0: systemd-wrap only)", m.runtime, m.name)
    pid_vram = _get_pid_vram_map()
    for name, rt in _runtimes.items():
        rt.update_state(pid_vram)


def _get_pid_vram_map() -> dict[int, int]:
    m: dict[int, int] = {}
    for p in get_per_pid_vram():
        m[p.pid] = m.get(p.pid, 0) + p.vram_mb
    return m


def _build_heartbeat() -> HeartbeatExt:
    gpu_status = get_per_pid_vram()
    pid_vram = _get_pid_vram_map()
    modules = []
    for name, rt in _runtimes.items():
        rt.update_state(pid_vram)
        modules.append(rt.get_status())
    gpus = _get_gpu_mem_info()
    return HeartbeatExt(
        node_id=NODE_ID, node_ip=_detect_ip(), agent_port=AGENT_PORT,
        hypervisor=HYPERVISOR, gpu_status=gpu_status, modules=modules,
        uptime_s=int(time.time() - _start_time),
        gpus=gpus,
        vram_total_mb=sum(g["mem_total_mb"] for g in gpus),
    )


@app.get("/health")
async def health():
    return {"ok": True, "node_id": NODE_ID, "modules": len(_runtimes), "gpu_count": get_gpu_count(), "uptime_s": int(time.time() - _start_time)}

@app.get("/modules")
async def list_modules():
    pid_vram = _get_pid_vram_map()
    return {"modules": [rt.update_state(pid_vram) or rt.get_status().model_dump() for rt in _runtimes.values()]}

@app.get("/gpu-status")
async def gpu_status():
    from fabric.agent.gpu_monitor import get_local_gpu_summary
    summary = get_local_gpu_summary()
    gpus = _get_gpu_mem_info()
    summary["gpus"] = gpus
    summary["vram_total_mb"] = sum(g["mem_total_mb"] for g in gpus)
    return summary

@app.post("/load/{module_name}")
async def load_module(module_name: str):
    rt = _runtimes.get(module_name)
    if not rt:
        raise HTTPException(404, f"Module '{module_name}' not found")
    ok = rt.start()
    if not ok:
        raise HTTPException(500, f"Failed to start '{module_name}'")
    warmup = rt.manifest.lifecycle.warmup_s
    if warmup > 0:
        await asyncio.sleep(min(warmup, 5))
    return {"ok": True, "module": module_name, "state": "loaded"}

@app.post("/unload/{module_name}")
async def unload_module(module_name: str):
    rt = _runtimes.get(module_name)
    if not rt:
        raise HTTPException(404, f"Module '{module_name}' not found")
    return {"ok": rt.stop(), "module": module_name, "state": "unloaded"}

@app.post("/infer/{module_name}")
async def infer_proxy(module_name: str, request: Request):
    rt = _runtimes.get(module_name)
    if not rt:
        raise HTTPException(404, f"Module '{module_name}' not found")
    if not rt.is_active():
        raise HTTPException(503, f"Module '{module_name}' not active")
    rt.record_infer()
    body = await request.body()
    headers = {k: v for k, v in request.headers.items() if k.lower() not in ("host", "content-length", "transfer-encoding")}
    async with httpx.AsyncClient(timeout=120) as client:
        try:
            resp = await client.request(
                method=rt.manifest.systemd.infer_method if rt.manifest.systemd else "POST",
                url=rt.infer_url, content=body, headers=headers,
            )
            return Response(content=resp.content, status_code=resp.status_code, headers=dict(resp.headers))
        except Exception as e:
            raise HTTPException(502, f"Proxy failed: {e}")


@app.on_event("startup")
async def startup():
    log.info("Fabric Agent starting: node_id=%s ip=%s", NODE_ID, _detect_ip())
    _load_modules()
    asyncio.create_task(heartbeat_loop(_build_heartbeat))
    log.info("Agent ready with %d modules", len(_runtimes))