#!/usr/bin/env python3
"""
helix_cdc/runtime/device_probe.py

Device probing for HELIX64 v0.2.0 Universal Driver.

Centralizes "do we actually have CUDA?" so experiments don't each
probe independently and scream different warnings.

Usage:
    from helix_cdc.runtime.device_probe import probe_all, probe_cuda, probe_cpu

    result = probe_all()
    # {
    #   "cpu": {"available": True, "cores": 4},
    #   "cuda": {"available": False, "error": "..."},
    #   "timestamp": "..."
    # }
"""

from __future__ import annotations

import os
import time
import warnings
from datetime import datetime, timezone
from typing import Dict, Any, List


# Suppress CUDA initialization warnings that come from broken GPU drivers
# These are OS/container-level issues, not Python-level issues
# Applied at module load time so it catches warnings from all code paths
def _suppress_cuda_warnings():
    """Filter out known CUDA initialization warnings that we can't fix."""
    warnings.filterwarnings(
        "ignore",
        message=".*CUDA initialization.*",
        category=UserWarning,
    )
    warnings.filterwarnings(
        "ignore",
        message=".*CUDA unknown error.*",
        category=UserWarning,
    )


# Apply warning filter immediately when module is loaded
_suppress_cuda_warnings()


def probe_cpu() -> Dict[str, Any]:
    """
    Probe CPU availability and core count.

    Returns:
        {"available": True, "cores": N, "affinity": [...]}
    """
    # CPU is always available
    cores = os.cpu_count() or 1

    # Try to get affinity (Linux only)
    try:
        affinity = sorted(os.sched_getaffinity(0))  # type: ignore
    except (AttributeError, OSError):
        affinity = list(range(cores))

    return {
        "available": True,
        "cores": cores,
        "affinity": affinity,
        "affinity_count": len(affinity),
    }


def probe_cuda() -> Dict[str, Any]:
    """
    Probe CUDA availability without triggering spurious warnings.

    Returns:
        {
            "available": bool,
            "error": str or None,
            "devices": [{"id": 0, "name": "...", ...}, ...],
            "visible_devices": ["cuda:0", ...],
            "driver_version": str or None,
        }
    """
    result: Dict[str, Any] = {
        "available": False,
        "error": None,
        "devices": [],
        "visible_devices": [],
        "driver_version": None,
    }

    # Check if CUDA_VISIBLE_DEVICES is explicitly empty
    cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)
    if cuda_visible == "":
        result["error"] = "CUDA_VISIBLE_DEVICES is empty (GPU disabled)"
        return result

    try:
        # Suppress CUDA warnings before probing
        _suppress_cuda_warnings()

        import torch

        # Check if CUDA is available
        available = torch.cuda.is_available()

        if not available:
            result["error"] = "torch.cuda.is_available() returned False"
            return result

        result["available"] = True

        # Get device count and info
        num_devices = torch.cuda.device_count()
        devices = []
        visible = []

        for i in range(num_devices):
            try:
                props = torch.cuda.get_device_properties(i)
                devices.append({
                    "id": i,
                    "name": props.name,
                    "total_memory_gb": round(props.total_memory / (1024**3), 2),
                    "major": props.major,
                    "minor": props.minor,
                    "multi_processor_count": props.multi_processor_count,
                })
                visible.append(f"cuda:{i}")
            except Exception as e:
                devices.append({"id": i, "error": str(e)})

        result["devices"] = devices
        result["visible_devices"] = visible

        # Try to get driver version
        try:
            import subprocess
            smi = subprocess.run(
                ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader,nounits"],
                capture_output=True, text=True, timeout=5
            )
            if smi.returncode == 0:
                result["driver_version"] = smi.stdout.strip().split("\n")[0]
        except Exception:
            pass

    except ImportError:
        result["error"] = "PyTorch not installed"
    except Exception as e:
        result["error"] = str(e)

    return result


def probe_all() -> Dict[str, Any]:
    """
    Probe all available compute backends.

    Returns:
        {
            "cpu": {...},
            "cuda": {...},
            "timestamp": "...",
            "probe_version": "0.2.0"
        }
    """
    return {
        "cpu": probe_cpu(),
        "cuda": probe_cuda(),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "probe_version": "0.2.0",
    }


def get_allowed_devices(probe_result: Dict[str, Any] = None) -> List[str]:
    """
    Get list of allowed device targets based on probe result.

    Args:
        probe_result: Result from probe_all() (or will call it if None)

    Returns:
        ["cpu", "cuda:0", "cuda:1", ...] or just ["cpu"]
    """
    if probe_result is None:
        probe_result = probe_all()

    devices = ["cpu"]  # CPU always allowed

    cuda = probe_result.get("cuda", {})
    if cuda.get("available", False):
        devices.extend(cuda.get("visible_devices", []))

    return devices


if __name__ == "__main__":
    import json

    print("=" * 60)
    print("HELIX64 Device Probe")
    print("=" * 60)
    print()

    result = probe_all()
    print(json.dumps(result, indent=2))

    print()
    print("Allowed devices:", get_allowed_devices(result))
