"""GPU detection and metrics for NVIDIA, AMD, and Apple Silicon GPUs."""

import base64
import json as _json
import logging
import os
import platform
import subprocess
from pathlib import Path
from typing import Optional

from models import GPUInfo, IndividualGPU

logger = logging.getLogger(__name__)


def run_command(cmd: list[str], timeout: int = 5) -> tuple[bool, str]:
    """Run a shell command and return (success, output)."""
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return result.returncode == 0, result.stdout.strip()
    except subprocess.TimeoutExpired:
        return False, "timeout"
    except (subprocess.SubprocessError, OSError) as e:
        return False, str(e)


def _read_sysfs(path: str) -> Optional[str]:
    """Read a sysfs file, returning None on failure."""
    try:
        with open(path, "r") as f:
            return f.read().strip()
    except (OSError, IOError):
        return None


def _read_meminfo_mb() -> Optional[tuple[int, int]]:
    """Return (used_mb, total_mb) from /proc/meminfo, or None on failure.

    Used as a fallback for unified-memory NVIDIA GPUs (GB10, GB200) where
    nvidia-smi reports [N/A] for memory.total/memory.used. Mirrors the
    installer's detection.sh fallback.
    """
    try:
        with open("/proc/meminfo") as f:
            info = {}
            for line in f:
                parts = line.split()
                if len(parts) >= 2:
                    info[parts[0].rstrip(":")] = int(parts[1])
    except OSError:
        return None
    total_kb = info.get("MemTotal", 0)
    avail_kb = info.get("MemAvailable", 0)
    if total_kb <= 0:
        return None
    used_mb = max(0, (total_kb - avail_kb)) // 1024
    total_mb = total_kb // 1024
    return used_mb, total_mb


def _find_amd_gpu_sysfs() -> Optional[str]:
    """Find the sysfs base path for an AMD GPU device."""
    import glob
    for card_dir in sorted(glob.glob("/sys/class/drm/card*/device")):
        vendor = _read_sysfs(f"{card_dir}/vendor")
        if vendor == "0x1002":
            return card_dir
    return None


def _find_hwmon_dir(device_path: str) -> Optional[str]:
    """Find the hwmon directory for an AMD GPU device."""
    import glob
    hwmon_dirs = sorted(glob.glob(f"{device_path}/hwmon/hwmon*"))
    return hwmon_dirs[0] if hwmon_dirs else None


def get_gpu_info_amd() -> Optional[GPUInfo]:
    """Get GPU metrics from amdgpu sysfs."""
    base = _find_amd_gpu_sysfs()
    if not base:
        return None

    hwmon = _find_hwmon_dir(base)

    try:
        vram_total_str = _read_sysfs(f"{base}/mem_info_vram_total")
        vram_used_str = _read_sysfs(f"{base}/mem_info_vram_used")
        gtt_total_str = _read_sysfs(f"{base}/mem_info_gtt_total")
        gtt_used_str = _read_sysfs(f"{base}/mem_info_gtt_used")
        gpu_busy_str = _read_sysfs(f"{base}/gpu_busy_percent")

        if not vram_total_str or not vram_used_str:
            return None

        vram_total = int(vram_total_str)
        vram_used = int(vram_used_str)
        gtt_total = int(gtt_total_str) if gtt_total_str else 0
        gtt_used = int(gtt_used_str) if gtt_used_str else 0
        gpu_busy = int(gpu_busy_str) if gpu_busy_str else 0

        is_unified = gtt_total > vram_total * 4

        if is_unified:
            mem_total = gtt_total
            mem_used = gtt_used
        else:
            mem_total = vram_total
            mem_used = vram_used

        temp = 0
        power_w = None
        if hwmon:
            temp_str = _read_sysfs(f"{hwmon}/temp1_input")
            if temp_str:
                temp = int(temp_str) // 1000

            power_str = _read_sysfs(f"{hwmon}/power1_average")
            if power_str:
                power_w = round(int(power_str) / 1e6, 1)

        gpu_name = _read_sysfs(f"{base}/product_name")
        memory_type = "unified" if is_unified else "discrete"
        if not gpu_name:
            if is_unified:
                gpu_name = get_gpu_tier(mem_total / (1024**3), memory_type)
            else:
                gpu_name = "AMD Radeon"

        mem_used_mb = mem_used // (1024 * 1024)
        mem_total_mb = mem_total // (1024 * 1024)

        return GPUInfo(
            name=gpu_name,
            memory_used_mb=mem_used_mb,
            memory_total_mb=mem_total_mb,
            memory_percent=round(mem_used_mb / mem_total_mb * 100, 1) if mem_total_mb > 0 else 0,
            utilization_percent=gpu_busy,
            temperature_c=temp,
            power_w=power_w,
            memory_type=memory_type,
            gpu_backend="amd",
        )
    except (ValueError, TypeError):
        return None


def get_gpu_info_nvidia() -> Optional[GPUInfo]:
    """Get GPU metrics from nvidia-smi.

    Handles multi-GPU systems by summing VRAM across all GPUs and
    reporting aggregate utilization and peak temperature.
    """
    success, output = run_command([
        "nvidia-smi",
        "--query-gpu=name,memory.used,memory.total,utilization.gpu,temperature.gpu,power.draw",
        "--format=csv,noheader,nounits"
    ])

    if not success or not output:
        return None

    # nvidia-smi returns one line per GPU; split before parsing
    lines = [line.strip() for line in output.strip().splitlines() if line.strip()]
    if not lines:
        return None

    try:
        gpus = []
        any_unified = False
        for line in lines:
            parts = [p.strip() for p in line.split(",")]
            if len(parts) < 5:
                continue
            power_w = None
            if len(parts) >= 6 and parts[5] not in ("[N/A]", "[Not Supported]", "N/A", "Not Supported", ""):
                try:
                    power_w = round(float(parts[5]), 1)
                except (ValueError, TypeError):
                    pass
            na_values = ("[N/A]", "[Not Supported]", "N/A", "Not Supported", "")
            # GB10/GB200 unified memory: nvidia-smi reports [N/A] for memory
            # fields; fall back to /proc/meminfo (mirrors detection.sh).
            unified = parts[1] in na_values or parts[2] in na_values
            if unified:
                fallback = _read_meminfo_mb()
                if not fallback:
                    continue
                mem_used, mem_total = fallback
                any_unified = True
            else:
                mem_used = int(parts[1])
                mem_total = int(parts[2])
            util = int(parts[3]) if parts[3] not in na_values else 0
            temp = int(parts[4]) if parts[4] not in na_values else 0
            gpus.append({
                "name": parts[0],
                "mem_used": mem_used,
                "mem_total": mem_total,
                "util": util,
                "temp": temp,
                "power_w": power_w,
            })

        if not gpus:
            return None

        memory_type = "unified" if any_unified else "discrete"

        if len(gpus) == 1:
            g = gpus[0]
            mem_used, mem_total = g["mem_used"], g["mem_total"]
            return GPUInfo(
                name=g["name"],
                memory_used_mb=mem_used,
                memory_total_mb=mem_total,
                memory_percent=round(mem_used / mem_total * 100, 1) if mem_total > 0 else 0,
                utilization_percent=g["util"],
                temperature_c=g["temp"],
                power_w=g["power_w"],
                memory_type=memory_type,
                gpu_backend="nvidia",
            )

        # Multi-GPU: aggregate across all GPUs
        mem_used = sum(g["mem_used"] for g in gpus)
        mem_total = sum(g["mem_total"] for g in gpus)
        avg_util = round(sum(g["util"] for g in gpus) / len(gpus))
        max_temp = max(g["temp"] for g in gpus)
        total_power: Optional[float] = None
        power_values = [g["power_w"] for g in gpus if g["power_w"] is not None]
        if power_values:
            total_power = round(sum(power_values), 1)

        # Build a display name: "RTX 4090 × 2" or "RTX 3090 + RTX 4090"
        names = [g["name"] for g in gpus]
        if len(set(names)) == 1:
            display_name = f"{names[0]} \u00d7 {len(gpus)}"
        else:
            display_name = " + ".join(names[:2])
            if len(names) > 2:
                display_name += f" + {len(names) - 2} more"

        return GPUInfo(
            name=display_name,
            memory_used_mb=mem_used,
            memory_total_mb=mem_total,
            memory_percent=round(mem_used / mem_total * 100, 1) if mem_total > 0 else 0,
            utilization_percent=avg_util,
            temperature_c=max_temp,
            power_w=total_power,
            memory_type=memory_type,
            gpu_backend="nvidia",
        )
    except (ValueError, IndexError):
        pass

    return None


def get_gpu_info_apple() -> Optional[GPUInfo]:
    """Get GPU metrics for Apple Silicon via system_profiler (native) or env vars (container)."""
    gpu_backend = os.environ.get("GPU_BACKEND", "").lower()

    if platform.system() == "Darwin":
        try:
            # Get chip name
            success, chip_output = run_command(["sysctl", "-n", "machdep.cpu.brand_string"])
            chip_name = chip_output.strip() if success else "Apple Silicon"

            # Get total memory (unified memory on Apple Silicon)
            success, mem_output = run_command(["sysctl", "-n", "hw.memsize"])
            if not success:
                return None

            total_bytes = int(mem_output.strip())
            total_mb = total_bytes // (1024 * 1024)

            # Estimate used memory from vm_stat
            used_mb = 0
            success, vm_output = run_command(["vm_stat"])
            if success:
                import re
                pages = {}
                for line in vm_output.splitlines():
                    match = re.match(r"(.+?):\s+(\d+)", line)
                    if match:
                        pages[match.group(1).strip()] = int(match.group(2))
                page_size = 16384
                ps_match = re.search(r"page size of (\d+) bytes", vm_output)
                if ps_match:
                    page_size = int(ps_match.group(1))
                active = pages.get("Pages active", 0)
                wired = pages.get("Pages wired down", 0)
                compressed = pages.get("Pages occupied by compressor", 0)
                used_mb = (active + wired + compressed) * page_size // (1024 * 1024)

            return GPUInfo(
                name=chip_name,
                memory_used_mb=used_mb,
                memory_total_mb=total_mb,
                memory_percent=round(used_mb / total_mb * 100, 1) if total_mb > 0 else 0,
                utilization_percent=0,  # not easily available without IOKit
                temperature_c=0,
                power_w=None,
                memory_type="unified",
                gpu_backend="apple",
            )
        except (ValueError, TypeError) as e:
            logger.debug("Apple Silicon GPU detection failed: %s", e)
            return None

    elif gpu_backend == "apple":
        # Linux container path (Docker Desktop on macOS): use HOST_RAM_GB env var
        host_ram_gb_str = os.environ.get("HOST_RAM_GB", "")
        if not host_ram_gb_str:
            return None
        try:
            host_ram_gb_float = float(host_ram_gb_str)
        except ValueError:
            return None
        if host_ram_gb_float <= 0:
            return None
        total_mb = int(host_ram_gb_float * 1024)
        # Use /proc/meminfo for used memory (best available proxy inside container)
        # Note: used_mb reflects Docker Desktop VM memory pressure, not the host Mac's.
        # Total is correctly overridden by HOST_RAM_GB. See issue #102 for a future
        # host-metrics collector that would fix used_mb.
        used_mb = 0
        try:
            with open("/proc/meminfo") as f:
                meminfo = {}
                for line in f:
                    parts = line.split()
                    if len(parts) >= 2:
                        meminfo[parts[0].rstrip(":")] = int(parts[1])
            avail = meminfo.get("MemAvailable", 0)
            total_kb = meminfo.get("MemTotal", 0)
            used_mb = (total_kb - avail) // 1024
        except OSError:
            pass
        return GPUInfo(
            name=f"Apple M-Series ({int(host_ram_gb_float)} GB Unified)",
            memory_used_mb=used_mb,
            memory_total_mb=total_mb,
            memory_percent=round(used_mb / total_mb * 100, 1) if total_mb > 0 else 0,
            utilization_percent=0,
            temperature_c=0,
            power_w=None,
            memory_type="unified",
            gpu_backend="apple",
        )

    return None


def get_gpu_info() -> Optional[GPUInfo]:
    """Get GPU metrics. Tries the configured backend first, then auto-detects."""
    gpu_backend = os.environ.get("GPU_BACKEND", "").lower()

    if gpu_backend == "amd":
        info = get_gpu_info_amd()
        if info:
            return info

    if gpu_backend == "apple":
        info = get_gpu_info_apple()
        if info:
            return info

    info = get_gpu_info_nvidia()
    if info:
        return info

    if gpu_backend != "amd":
        info = get_gpu_info_amd()
        if info:
            return info

    # Auto-detect Apple Silicon if no backend specified and nothing else found
    if platform.system() == "Darwin":
        return get_gpu_info_apple()

    return None


# ============================================================================
# Topology — read from file written by installer / dream-cli
# ============================================================================

def read_gpu_topology() -> Optional[dict]:
    """Read GPU topology from config/gpu-topology.json if it exists.

    The file is written by the installer (03-features.sh) and refreshed by
    'dream gpu reassign'.  Inside the API container it is available at
    /dream-server/config/gpu-topology.json (mounted read-only).
    """
    install_dir = os.environ.get("DREAM_INSTALL_DIR", os.path.expanduser("~/dream-server"))
    topo_path = Path(install_dir) / "config" / "gpu-topology.json"
    if not topo_path.exists():
        logger.warning("Topology file not found at %s", topo_path)
        return None
    try:
        return _json.loads(topo_path.read_text())
    except (OSError, _json.JSONDecodeError) as exc:
        logger.warning("Failed to read topology file %s: %s", topo_path, exc)
        return None


# ============================================================================
# Assignment decoding helpers
# ============================================================================

def decode_gpu_assignment() -> Optional[dict]:
    """Decode GPU_ASSIGNMENT_JSON_B64, preferring the live .env file over the
    container startup environment so reassignments are reflected without restart."""
    b64 = _read_env_var_from_file("GPU_ASSIGNMENT_JSON_B64") or os.environ.get("GPU_ASSIGNMENT_JSON_B64", "")
    if not b64:
        return None
    try:
        return _json.loads(base64.b64decode(b64.strip()).decode("utf-8"))
    except (base64.binascii.Error, _json.JSONDecodeError, UnicodeDecodeError):
        return None


def _read_env_var_from_file(key: str) -> str:
    """Read a single variable directly from the .env file (split on first '=' only)."""
    install_dir = os.environ.get("DREAM_INSTALL_DIR", os.path.expanduser("~/dream-server"))
    env_path = Path(install_dir) / ".env"
    try:
        for line in env_path.read_text().splitlines():
            if line.startswith(f"{key}="):
                return line[len(key) + 1:].strip().strip("\"'")
    except OSError:
        pass
    return ""


def _build_uuid_service_map(assignment: dict) -> dict[str, list[str]]:
    """Map GPU UUID → list of service names from assignment JSON."""
    result: dict[str, list[str]] = {}
    services = assignment.get("gpu_assignment", {}).get("services", {})
    for svc_name, svc_data in services.items():
        for uuid in svc_data.get("gpus", []):
            result.setdefault(uuid, []).append(svc_name)
    return result


def _infer_gpu_services_from_processes() -> dict[str, list[str]]:
    """Fallback: infer GPU service assignment from nvidia-smi compute processes.

    Used when GPU_ASSIGNMENT_JSON_B64 is not set (single-GPU setups).
    Maps GPU UUID -> list of likely service names based on running processes.
    """
    success, output = run_command([
        "nvidia-smi",
        "--query-compute-apps=gpu_uuid,pid,used_memory",
        "--format=csv,noheader,nounits",
    ])
    if not success or not output:
        return {}

    # Collect UUIDs that have active compute processes
    active_uuids: dict[str, int] = {}  # uuid -> total used memory MB
    for line in output.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) >= 3:
            uuid = parts[0]
            try:
                mem = int(parts[2])
            except ValueError:
                mem = 0
            active_uuids[uuid] = active_uuids.get(uuid, 0) + mem

    if not active_uuids:
        return {}

    # For each active GPU, attribute to llama-server (the primary GPU consumer)
    result: dict[str, list[str]] = {}
    for uuid, mem_mb in active_uuids.items():
        if mem_mb > 100:
            result[uuid] = ["llama-server"]

    return result


# ============================================================================
# Per-GPU detailed detection
# ============================================================================

def get_gpu_info_nvidia_detailed() -> Optional[list[IndividualGPU]]:
    """Return one IndividualGPU per NVIDIA GPU, with assigned_services populated.

    Returns None if nvidia-smi is unavailable or returns no data.
    """
    success, output = run_command([
        "nvidia-smi",
        "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu,temperature.gpu,power.draw",
        "--format=csv,noheader,nounits",
    ])
    if not success or not output:
        return None

    lines = [ln.strip() for ln in output.strip().splitlines() if ln.strip()]
    if not lines:
        return None

    assignment = decode_gpu_assignment()
    if assignment:
        uuid_service_map = _build_uuid_service_map(assignment)
    else:
        uuid_service_map = _infer_gpu_services_from_processes()

    gpus: list[IndividualGPU] = []
    na_values = ("[N/A]", "[Not Supported]", "N/A", "Not Supported", "")
    for line in lines:
        try:
            parts = [p.strip() for p in line.split(",")]
            if len(parts) < 7:
                continue
            power_w = None
            if len(parts) >= 8 and parts[7] not in na_values:
                try:
                    power_w = round(float(parts[7]), 1)
                except (ValueError, TypeError):
                    pass
            # GB10/GB200 unified memory fallback (see _read_meminfo_mb).
            if parts[3] in na_values or parts[4] in na_values:
                fallback = _read_meminfo_mb()
                if not fallback:
                    continue
                mem_used, mem_total = fallback
            else:
                mem_used = int(parts[3])
                mem_total = int(parts[4])
            uuid = parts[1]
            gpus.append(IndividualGPU(
                index=int(parts[0]),
                uuid=uuid,
                name=parts[2],
                memory_used_mb=mem_used,
                memory_total_mb=mem_total,
                memory_percent=round(mem_used / mem_total * 100, 1) if mem_total > 0 else 0.0,
                utilization_percent=int(parts[5]),
                temperature_c=int(parts[6]),
                power_w=power_w,
                assigned_services=uuid_service_map.get(uuid, []),
            ))
        except (ValueError, IndexError):
            logger.warning("Skipping unparseable nvidia-smi row: %s", line)

    return gpus or None


def get_gpu_info_amd_detailed() -> Optional[list[IndividualGPU]]:
    """Return one IndividualGPU per AMD GPU by iterating all amdgpu sysfs cards.

    Returns None if no AMD GPUs are found.
    """
    import glob as _glob
    card_dirs = sorted(_glob.glob("/sys/class/drm/card*/device"))
    amd_cards = [d for d in card_dirs if _read_sysfs(f"{d}/vendor") == "0x1002"]
    if not amd_cards:
        return None

    gpus: list[IndividualGPU] = []
    for idx, base in enumerate(amd_cards):
        hwmon = _find_hwmon_dir(base)
        try:
            vram_total_str = _read_sysfs(f"{base}/mem_info_vram_total")
            vram_used_str = _read_sysfs(f"{base}/mem_info_vram_used")
            gtt_total_str = _read_sysfs(f"{base}/mem_info_gtt_total")
            gtt_used_str = _read_sysfs(f"{base}/mem_info_gtt_used")
            gpu_busy_str = _read_sysfs(f"{base}/gpu_busy_percent")

            if not vram_total_str or not vram_used_str:
                continue

            vram_total = int(vram_total_str)
            vram_used = int(vram_used_str)
            gtt_total = int(gtt_total_str) if gtt_total_str else 0
            gtt_used = int(gtt_used_str) if gtt_used_str else 0
            gpu_busy = int(gpu_busy_str) if gpu_busy_str else 0

            is_unified = gtt_total > vram_total * 4
            mem_total = gtt_total if is_unified else vram_total
            mem_used = gtt_used if is_unified else vram_used

            temp = 0
            power_w = None
            if hwmon:
                temp_str = _read_sysfs(f"{hwmon}/temp1_input")
                if temp_str:
                    temp = int(temp_str) // 1000
                power_str = _read_sysfs(f"{hwmon}/power1_average")
                if power_str:
                    power_w = round(int(power_str) / 1e6, 1)

            gpu_name = _read_sysfs(f"{base}/product_name") or "AMD Radeon"
            card_id = base.split("/")[-2]  # "card0", "card1", …
            mem_used_mb = mem_used // (1024 * 1024)
            mem_total_mb = mem_total // (1024 * 1024)

            gpus.append(IndividualGPU(
                index=idx,
                uuid=card_id,
                name=gpu_name,
                memory_used_mb=mem_used_mb,
                memory_total_mb=mem_total_mb,
                memory_percent=round(mem_used_mb / mem_total_mb * 100, 1) if mem_total_mb > 0 else 0.0,
                utilization_percent=gpu_busy,
                temperature_c=temp,
                power_w=power_w,
                assigned_services=[],  # AMD assignment uses NVIDIA UUIDs; not mapped here
            ))
        except (ValueError, TypeError):
            continue

    return gpus or None


def get_gpu_tier(vram_gb: float, memory_type: str = "discrete") -> str:
    """Get tier name based on VRAM."""
    if memory_type == "unified":
        if vram_gb >= 90:
            return "Strix Halo 90+"
        else:
            return "Strix Halo Compact"
    if vram_gb >= 80:
        return "Professional"
    elif vram_gb >= 24:
        return "Prosumer"
    elif vram_gb >= 16:
        return "Standard"
    elif vram_gb >= 8:
        return "Entry"
    else:
        return "Minimal"
