"""GPU observation via nvidia-smi — collects utilization, memory, thermal, power."""

from __future__ import annotations

import asyncio
import json
import logging
import subprocess
from datetime import datetime, timedelta, timezone
from typing import Optional

from radix_edge.config import get_config
from radix_edge.db import get_db

logger = logging.getLogger("radix_edge.observer")

# nvidia-smi query fields
_QUERY_FIELDS = (
    "index,uuid,name,memory.total,memory.used,memory.free,"
    "utilization.gpu,temperature.gpu,power.draw,power.limit,"
    "pcie.link.gen.current"
)


def _parse_mib(val: str) -> int:
    """Parse nvidia-smi MiB value like '81920 MiB' -> 81920."""
    return int(val.strip().split()[0]) if val.strip() and val.strip() != "[N/A]" else 0


def _parse_float(val: str) -> float:
    """Parse nvidia-smi float value like '45.2 %' or '250.00 W' -> float."""
    val = val.strip()
    if not val or val == "[N/A]":
        return 0.0
    # Remove units like %, W, etc.
    parts = val.split()
    try:
        return float(parts[0])
    except (ValueError, IndexError):
        return 0.0


def poll_gpus() -> list[dict]:
    """Run nvidia-smi and return parsed GPU data. Returns empty list if nvidia-smi unavailable."""
    try:
        result = subprocess.run(
            ["nvidia-smi", f"--query-gpu={_QUERY_FIELDS}", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            logger.warning("nvidia-smi returned code %d: %s", result.returncode, result.stderr.strip())
            return []
    except FileNotFoundError:
        logger.info("nvidia-smi not found — no GPUs detected")
        return []
    except subprocess.TimeoutExpired:
        logger.warning("nvidia-smi timed out")
        return []

    gpus = []
    for line in result.stdout.strip().split("\n"):
        if not line.strip():
            continue
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 11:
            continue
        gpus.append({
            "gpu_index": int(parts[0]),
            "uuid": parts[1],
            "name": parts[2],
            "memory_total_mb": _parse_mib(parts[3]),
            "memory_used_mb": _parse_mib(parts[4]),
            "memory_free_mb": _parse_mib(parts[5]),
            "utilization_pct": _parse_float(parts[6]),
            "temperature_c": _parse_float(parts[7]),
            "power_draw_w": _parse_float(parts[8]),
            "power_limit_w": _parse_float(parts[9]),
            "pcie_gen": int(_parse_float(parts[10])),
        })
    return gpus


def store_gpu_state(db_conn, gpus: list[dict], topology: Optional[dict[int, list[int]]] = None):
    """Upsert GPU state into the gpus table and append to gpu_observations."""
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")

    for gpu in gpus:
        idx = gpu["gpu_index"]
        nvlink_peers = json.dumps(topology.get(idx, [])) if topology else "[]"

        # Preserve assigned_job_id from existing row
        existing = db_conn.execute("SELECT assigned_job_id, status FROM gpus WHERE gpu_index = ?", (idx,)).fetchone()
        assigned_job_id = existing["assigned_job_id"] if existing else None
        status = existing["status"] if existing and existing["status"] != "idle" else "idle"

        db_conn.execute(
            """INSERT INTO gpus (gpu_index, uuid, name, memory_total_mb, memory_used_mb, memory_free_mb,
                utilization_pct, temperature_c, power_draw_w, power_limit_w, pcie_gen,
                nvlink_peers, assigned_job_id, status, last_seen_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(gpu_index) DO UPDATE SET
                uuid=excluded.uuid, name=excluded.name,
                memory_total_mb=excluded.memory_total_mb, memory_used_mb=excluded.memory_used_mb,
                memory_free_mb=excluded.memory_free_mb, utilization_pct=excluded.utilization_pct,
                temperature_c=excluded.temperature_c, power_draw_w=excluded.power_draw_w,
                power_limit_w=excluded.power_limit_w, pcie_gen=excluded.pcie_gen,
                nvlink_peers=excluded.nvlink_peers,
                assigned_job_id=?, status=?, last_seen_at=?""",
            (
                idx, gpu["uuid"], gpu["name"], gpu["memory_total_mb"], gpu["memory_used_mb"],
                gpu["memory_free_mb"], gpu["utilization_pct"], gpu["temperature_c"],
                gpu["power_draw_w"], gpu["power_limit_w"], gpu["pcie_gen"],
                nvlink_peers, assigned_job_id, status, now,
                # ON CONFLICT SET values:
                assigned_job_id, status, now,
            ),
        )

        # Append observation history
        db_conn.execute(
            """INSERT INTO gpu_observations (gpu_index, utilization_pct, memory_used_mb, temperature_c, power_draw_w, observed_at)
            VALUES (?, ?, ?, ?, ?, ?)""",
            (idx, gpu["utilization_pct"], gpu["memory_used_mb"], gpu["temperature_c"], gpu["power_draw_w"], now),
        )

    db_conn.commit()


def prune_old_observations(db_conn, retention_hours: int = 24):
    """Remove gpu_observations older than retention window."""
    cutoff = (datetime.now(timezone.utc) - timedelta(hours=retention_hours)).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
    db_conn.execute("DELETE FROM gpu_observations WHERE observed_at < ?", (cutoff,))
    db_conn.commit()


def compute_observation_vector(db_conn) -> dict[str, float]:
    """Compute the 11-dimensional observation vector from local state for FEP.

    Returns dict with keys matching the observation dimensions.
    """
    # GPU utilization stats
    gpu_rows = db_conn.execute("SELECT utilization_pct, memory_used_mb, memory_total_mb, temperature_c, power_draw_w, power_limit_w, nvlink_peers, status FROM gpus").fetchall()

    if not gpu_rows:
        return {k: 0.0 for k in [
            "gpu_util_mean", "gpu_util_std", "queue_depth", "queue_depth_max",
            "avg_wait_seconds", "error_rate", "throughput_jobs_per_min",
            "memory_pressure", "thermal_headroom", "power_headroom", "topology_utilization",
        ]}

    utils = [r["utilization_pct"] for r in gpu_rows]
    n = len(utils)
    gpu_util_mean = sum(utils) / n if n else 0.0
    gpu_util_std = (sum((u - gpu_util_mean) ** 2 for u in utils) / n) ** 0.5 if n > 1 else 0.0

    # Queue depth
    queue_depth = db_conn.execute("SELECT COUNT(*) as c FROM jobs WHERE status = 'queued'").fetchone()["c"]

    # Max queue depth in last 5 minutes (approximate from current)
    queue_depth_max = float(queue_depth)

    # Average wait time for recently completed jobs (last hour)
    wait_row = db_conn.execute(
        """SELECT AVG(CAST((julianday(scheduled_at) - julianday(created_at)) * 86400 AS REAL)) as avg_wait
        FROM jobs WHERE scheduled_at IS NOT NULL
        AND created_at > datetime('now', '-1 hour')"""
    ).fetchone()
    avg_wait_seconds = wait_row["avg_wait"] or 0.0

    # Error rate (last hour)
    totals = db_conn.execute(
        """SELECT
            SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) as failed,
            SUM(CASE WHEN status IN ('succeeded', 'failed') THEN 1 ELSE 0 END) as total
        FROM jobs WHERE completed_at > datetime('now', '-1 hour')"""
    ).fetchone()
    error_rate = (totals["failed"] / totals["total"]) if totals["total"] and totals["total"] > 0 else 0.0

    # Throughput (jobs per minute, last hour)
    completed = db_conn.execute(
        "SELECT COUNT(*) as c FROM jobs WHERE status = 'succeeded' AND completed_at > datetime('now', '-1 hour')"
    ).fetchone()["c"]
    throughput_jobs_per_min = completed / 60.0

    # Memory pressure
    mem_pressures = [r["memory_used_mb"] / r["memory_total_mb"] if r["memory_total_mb"] > 0 else 0.0 for r in gpu_rows]
    memory_pressure = sum(mem_pressures) / len(mem_pressures) if mem_pressures else 0.0

    # Thermal headroom (higher = worse, closer to limits)
    thermal_vals = [r["temperature_c"] / 90.0 for r in gpu_rows]  # 90C typical thermal limit
    thermal_headroom = sum(thermal_vals) / len(thermal_vals) if thermal_vals else 0.0

    # Power headroom
    power_vals = [r["power_draw_w"] / r["power_limit_w"] if r["power_limit_w"] > 0 else 0.0 for r in gpu_rows]
    power_headroom = sum(power_vals) / len(power_vals) if power_vals else 0.0

    # Topology utilization: fraction of NVLink-connected GPU pairs where both are busy
    busy_indices = {r["gpu_index"] for r in db_conn.execute("SELECT gpu_index FROM gpus WHERE status = 'allocated'").fetchall()}
    nvlink_pairs = 0
    both_busy = 0
    for r in gpu_rows:
        peers = json.loads(r["nvlink_peers"]) if r["nvlink_peers"] else []
        for peer in peers:
            nvlink_pairs += 1
            if r["gpu_index"] in busy_indices and peer in busy_indices:
                both_busy += 1
    topology_utilization = both_busy / nvlink_pairs if nvlink_pairs > 0 else 0.0

    return {
        "gpu_util_mean": gpu_util_mean,
        "gpu_util_std": gpu_util_std,
        "queue_depth": float(queue_depth),
        "queue_depth_max": queue_depth_max,
        "avg_wait_seconds": avg_wait_seconds,
        "error_rate": error_rate,
        "throughput_jobs_per_min": throughput_jobs_per_min,
        "memory_pressure": memory_pressure,
        "thermal_headroom": thermal_headroom,
        "power_headroom": power_headroom,
        "topology_utilization": topology_utilization,
    }


async def observer_loop():
    """Background task: poll GPUs and update state continuously."""
    config = get_config()
    from radix_edge.observer.gpu_topology import detect_topology

    # Detect topology once at startup
    topology = detect_topology()
    if topology:
        logger.info("NVLink topology detected: %s", topology)
    else:
        logger.info("No NVLink topology detected (or nvidia-smi unavailable)")

    cycle = 0
    while True:
        try:
            gpus = poll_gpus()
            if gpus:
                with get_db() as db:
                    store_gpu_state(db, gpus, topology)

                    # Prune old observations every 100 cycles
                    if cycle % 100 == 0:
                        prune_old_observations(db, config.observer_retention_hours)

                if cycle == 0:
                    logger.info("Observer: %d GPUs detected", len(gpus))
            cycle += 1
        except Exception:
            logger.exception("Observer error")

        await asyncio.sleep(config.observer_interval)
