#!/usr/bin/env python3
"""GPU Fleet Monitor — Comprehensive monitoring server.

Monitors every subsystem in the inference harness:
  - GPU sidecars (VRAM, temp, util, power per card)
  - Router (model routing, circuit breakers, queue health)
  - LiteLLM (proxy health, key count, model sync)
  - Strix Halo (llama-server status, CPU load, context)
  - Dashboard (CT 116 harness-dashboard)

Endpoints:
  /          → GPU fleet HTML dashboard
  /gpu-data  → JSON GPU metrics (for dashboard API)
  /health    → Monitor self-health check

Run: python3 gpu-monitor-server.py
Default port: 9100
"""

import json, subprocess, http.server, threading, time, os
from datetime import datetime, timezone

DASHBOARD_PATH = "/root/dashboard/gpu-fleet.html"
PORT = 9100

# Cached data, refreshed every 15s
cache = {}
cache_lock = threading.Lock()

# Alert thresholds
THRESHOLDS = {
    "temp_c": {"warning": 80, "critical": 90},
    "vram_pct": {"warning": 90, "critical": 95},
    "gpu_util_pct": {"warning": 95, "critical": 98},
}


def http_get(url, timeout=5):
    """HTTP GET with timeout, returns parsed JSON or error dict."""
    import urllib.request
    try:
        resp = urllib.request.urlopen(url, timeout=timeout)
        return json.loads(resp.read())
    except Exception as e:
        return {"error": str(e)}


def poll_sidecar(host, port=8090):
    """Poll a GPU sidecar for raw metrics.
    Falls back to SSH-based nvidia-smi if sidecar is unreachable."""
    result = http_get(f"http://{host}:{port}/health", timeout=5)
    if "error" not in result:
        return result
    # Fallback: poll via SSH (explicit key path for non-interactive environments)
    try:
        proc = subprocess.run(
            ["/usr/bin/ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5",
             "-i", "/root/.ssh/id_ed25519", host,
             "nvidia-smi", "--query-gpu=name,temperature.gpu,utilization.gpu,"
             "utilization.memory,memory.used,memory.total,power.draw,power.limit,fan.speed",
             "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10
        )
        if proc.returncode == 0:
            parts = [p.strip() for p in proc.stdout.strip().split(",")]
            if len(parts) >= 8:
                def cv(v):
                    try: return float(v.replace("MiB","").replace("W","").replace("%","").strip())
                    except: return 0.0
                def mv(v):
                    try: return int(v.replace("MiB","").strip())
                    except: return 0
                return {
                    "gpu_name": parts[0], "temp_c": cv(parts[1]),
                    "gpu_util_pct": cv(parts[2]), "mem_util_pct": cv(parts[3]),
                    "vram_used_mb": mv(parts[4]), "vram_total_mb": mv(parts[5]),
                    "power_w": cv(parts[6]), "power_limit_w": cv(parts[7]),
                    "fan_pct": cv(parts[8])
                }
    except Exception as e:
        return {"error": str(e)}
    return result


def poll_router():
    """Poll router unified health through nginx (port 80)."""
    return http_get("http://192.168.68.116/health/unified", timeout=5)


def poll_router_health():
    """Poll router basic health through nginx."""
    return http_get("http://192.168.68.116/health", timeout=5)


def poll_litellm_health():
    """Poll LiteLLM — checks reachability via nginx proxy.
    
    LiteLLM health endpoints require authentication. We check if the
    service responds at all (even 401 = service is up). Also try the
    /key/list endpoint which confirms the proxy is fully functional.
    """
    import urllib.request
    import urllib.error
    urls_to_try = [
        "http://192.168.68.116/litellm/health",
        "http://192.168.68.116/litellm/ui/",
    ]
    last_error = None
    for url in urls_to_try:
        try:
            resp = urllib.request.urlopen(url, timeout=5)
            body = resp.read().decode()[:500]
            # Any response (even 401) means the proxy is routing to LiteLLM
            return {"reachable": True, "status_code": resp.getcode(),
                    "endpoint": url, "body_preview": body[:200]}
        except urllib.error.HTTPError as e:
            # HTTP error means the service IS reachable but returned error
            return {"reachable": True, "status_code": e.code,
                    "endpoint": url, "note": f"HTTP {e.code} (requires auth)"}
        except Exception as e:
            last_error = str(e)
    return {"error": last_error or "all endpoints failed"}


def poll_strix():
    """Poll Strix Halo — process check + CPU load + llama-server health."""
    try:
        # Check llama-server process
        proc = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5",
             "192.168.68.15",
             "pgrep -f llama-server > /dev/null 2>&1 && echo running || echo stopped"],
            capture_output=True, text=True, timeout=10
        )
        status = proc.stdout.strip()

        # Get CPU load
        load = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5",
             "192.168.68.15",
             "uptime | awk -F'load average:' '{print $2}' | tr -d ' '"],
            capture_output=True, text=True, timeout=10
        )
        cpu_load = load.stdout.strip() if load.returncode == 0 else "unknown"

        # Try llama-server health endpoint
        llama_health = http_get("http://192.168.68.15:8080/health", timeout=3)

        return {
            "status": status if status else "unknown",
            "cpu_load": cpu_load,
            "llama_health": llama_health if "error" not in llama_health else None
        }
    except Exception as e:
        return {"status": "error", "error": str(e)}


def poll_dashboard():
    """Check if harness-dashboard on CT 116 is serving."""
    return http_get("http://192.168.68.116/dashboard/", timeout=5)


def check_alerts(gpu_data):
    """Check GPU metrics against thresholds, return alert list."""
    alerts = []
    name = gpu_data.get("name", "unknown")
    for metric, thresholds in THRESHOLDS.items():
        # Map metric names from sidecar/router formats
        value = None
        if metric == "vram_pct":
            # Compute from vram_used_mb / vram_total_mb
            used = gpu_data.get("vram_used_mb", 0)
            total = gpu_data.get("vram_total_mb", 1)
            value = (used / total) * 100
        elif metric == "gpu_util_pct":
            value = gpu_data.get("gpu_util_pct", 0)
        elif metric == "temp_c":
            value = gpu_data.get("temp_c", 0)

        if value is not None:
            if value >= thresholds["critical"]:
                alerts.append({"gpu": name, "metric": metric, "level": "critical",
                               "value": round(value, 1), "threshold": thresholds["critical"]})
            elif value >= thresholds["warning"]:
                alerts.append({"gpu": name, "metric": metric, "level": "warning",
                               "value": round(value, 1), "threshold": thresholds["warning"]})
    return alerts


def compute_summary(router_data, gpus, litellm_data, strix_data):
    """Compute fleet-wide health summary."""
    # Count models available
    models_available = 0
    models_total = 0
    if "available_models" in router_data:
        models_available = len(router_data["available_models"])
        models_total = models_available  # router reports all expected

    # Count circuit breakers open
    cb_open = 0
    if "circuit_breaker" in router_data:
        for model, cb in router_data["circuit_breaker"].items():
            if cb.get("open", 0):
                cb_open += 1

    # Count GPU errors
    gpu_errors = sum(1 for g in gpus if "error" in g)

    # Fleet status
    if gpu_errors > 0 or cb_open > 0:
        fleet_status = "degraded"
    elif not router_data or "error" in router_data:
        fleet_status = "degraded"
    else:
        fleet_status = "healthy"

    litellm_ok = ("error" not in litellm_data) or litellm_data.get("reachable", False)
    return {
        "fleet_status": fleet_status,
        "models_available": models_available,
        "models_total": models_total,
        "circuit_breakers_open": cb_open,
        "gpu_count": len(gpus),
        "gpu_errors": gpu_errors,
        "strix_running": strix_data.get("status") == "running",
        "router_reachable": "error" not in router_data,
        "litellm_reachable": litellm_ok,
    }


def refresh_cache():
    """Refresh all fleet data from every subsystem."""
    global cache

    # 1. GPU sidecars
    gpu_map = {
        "192.168.68.8": {"name": "Dense (RTX 3090)", "models": ["gpu-dense"], "hostname": "ct8"},
        "192.168.68.110": {"name": "Light/Vision (RTX 5070)", "models": ["gpu-vision"], "hostname": "ct110"},
    }

    gpus = []
    all_alerts = []
    for host, info in gpu_map.items():
        data = poll_sidecar(host)
        if "error" not in data:
            data["name"] = info["name"]
            data["models"] = info["models"]
            data["hostname"] = info["hostname"]
            all_alerts.extend(check_alerts(data))
        else:
            data = {"name": info["name"], "hostname": info["hostname"], "error": data["error"]}
            all_alerts.append({"gpu": info["name"], "metric": "connectivity",
                               "level": "critical", "value": data["error"]})
        gpus.append(data)

    # 2. Router
    router = poll_router()
    router_basic = poll_router_health()
    router["_basic"] = router_basic

    # 3. LiteLLM
    litellm = poll_litellm_health()

    # 4. Strix Halo
    strix = poll_strix()

    # 5. Dashboard
    dashboard = poll_dashboard()

    # 6. Send Zulip DM for critical alerts (debounced 5 min)
    critical = [a for a in all_alerts if a.get("level") == "critical"]
    if critical and _debounce_alert():
        msg = "🔴 **GPU Fleet Alert**\n"
        for a in critical:
            msg += f"• {a['gpu']}: {a['metric']} = {a.get('value', a.get('actual', '?'))}\n"
        _send_zulip_dm(msg)

    # 7. Summary
    summary = compute_summary(router, gpus, litellm, strix)

    with cache_lock:
        cache = {
            "gpus": gpus,
            "strix": strix,
            "router": router,
            "litellm": litellm,
            "dashboard": {"reachable": "error" not in dashboard},
            "summary": summary,
            "alerts": all_alerts,
            "updated": datetime.now(timezone.utc).isoformat(),
            "monitor_version": "2.0.0",
        }


_last_alert_time = 0
_ALERT_COOLDOWN = 300  # 5 min between DMs

def _send_zulip_dm(msg):
    """Send a DM to the owner via Zulip."""
    try:
        site = "https://chat.sysloggh.net"
        email = "abiba-bot@chat.sysloggh.net"
        key = "cKTDMZAPW08dk3zl05sStzO7HRztzyn8"
        owner = 9
        auth_b64 = base64.b64encode(f"{email}:{key}".encode()).decode()
        data = urllib.parse.urlencode({"type": "private", "to": f"[{owner}]", "content": msg}).encode()
        req = urllib.request.Request(f"{site}/api/v1/messages", data=data,
            headers={"Authorization": f"Basic {auth_b64}"}, method="POST")
        urllib.request.urlopen(req, timeout=5)
    except Exception as e:
        print(f"[alert] DM failed: {e}")

def _debounce_alert():
    """Returns True if enough time has passed since last alert."""
    global _last_alert_time
    now = time.time()
    if now - _last_alert_time > _ALERT_COOLDOWN:
        _last_alert_time = now
        return True
    return False

class DashboardHandler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/gpu-data":
            with cache_lock:
                data = json.dumps(cache, indent=2)
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Access-Control-Allow-Origin", "*")
            self.end_headers()
            self.wfile.write(data.encode())
        elif self.path == "/health":
            with cache_lock:
                healthy = bool(cache and "summary" in cache)
            status = {"status": "healthy" if healthy else "starting",
                      "updated": cache.get("updated", "never") if cache else "never"}
            code = 200 if healthy else 503
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps(status).encode())
        elif self.path in ("/", "/index.html"):
            try:
                with open(DASHBOARD_PATH) as f:
                    html = f.read()
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(html.encode())
            except FileNotFoundError:
                self.send_response(404)
                self.end_headers()
                self.wfile.write(b"Dashboard not found")
        else:
            self.send_response(404)
            self.end_headers()

    def log_message(self, format, *args):
        pass  # Suppress request logs


def background_refresh():
    """Refresh data every 15 seconds."""
    while True:
        try:
            refresh_cache()
        except Exception as e:
            print(f"[WARN] Refresh failed: {e}", flush=True)
        time.sleep(15)


if __name__ == "__main__":
    print(f"Starting Comprehensive GPU Fleet Monitor v2.0.0 on port {PORT}...")
    print(f"  Dashboard: http://localhost:{PORT}/")
    print(f"  API:       http://localhost:{PORT}/gpu-data")
    print(f"  Health:    http://localhost:{PORT}/health")
    print(f"  Polling:   router (nginx:80), sidecars (:8090), strix, litellm, dashboard")

    # Initial fetch
    refresh_cache()

    # Start background refresher
    t = threading.Thread(target=background_refresh, daemon=True)
    t.start()

    # Start HTTP server
    server = http.server.HTTPServer(("0.0.0.0", PORT), DashboardHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")