Get Available Resources

#!/usr/bin/env python3
"""
System Resource Detection Script

Detects available compute resources including CPU, GPU, memory, and disk space.
Outputs a JSON file that Claude Code can use to make informed decisions about
computational approaches (e.g., whether to use Dask, Zarr, Joblib, etc.).

Supports: macOS, Linux, Windows
GPU Detection: NVIDIA (CUDA), AMD (ROCm), Apple Silicon (Metal)
"""

import datetime
import json
import os
import platform
import psutil
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Any, Optional


def get_cpu_info() -> Dict[str, Any]:
    """Detect CPU information."""
    cpu_info = {
        "physical_cores": psutil.cpu_count(logical=False),
        "logical_cores": psutil.cpu_count(logical=True),
        "max_frequency_mhz": None,
        "architecture": platform.machine(),
        "processor": platform.processor(),
    }

    # Get CPU frequency if available
    try:
        freq = psutil.cpu_freq()
        if freq:
            cpu_info["max_frequency_mhz"] = freq.max
            cpu_info["current_frequency_mhz"] = freq.current
    except Exception:
        pass

    return cpu_info


def get_memory_info() -> Dict[str, Any]:
    """Detect memory information."""
    mem = psutil.virtual_memory()
    swap = psutil.swap_memory()

    return {
        "total_gb": round(mem.total / (1024**3), 2),
        "available_gb": round(mem.available / (1024**3), 2),
        "used_gb": round(mem.used / (1024**3), 2),
        "percent_used": mem.percent,
        "swap_total_gb": round(swap.total / (1024**3), 2),
        "swap_available_gb": round((swap.total - swap.used) / (1024**3), 2),
    }


def get_disk_info(path: str = None) -> Dict[str, Any]:
    """Detect disk space information for working directory or specified path."""
    if path is None:
        path = os.getcwd()

    try:
        disk = psutil.disk_usage(path)
        return {
            "path": path,
            "total_gb": round(disk.total / (1024**3), 2),
            "available_gb": round(disk.free / (1024**3), 2),
            "used_gb": round(disk.used / (1024**3), 2),
            "percent_used": disk.percent,
        }
    except Exception as e:
        return {
            "path": path,
            "error": str(e),
        }


def detect_nvidia_gpus() -> List[Dict[str, Any]]:
    """Detect NVIDIA GPUs using nvidia-smi."""
    gpus = []

    try:
        # Try to run nvidia-smi
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,name,memory.total,memory.free,driver_version,compute_cap",
             "--format=csv,noheader,nounits"],
            capture_output=True,
            text=True,
            timeout=5
        )

        if result.returncode == 0:
            for line in result.stdout.strip().split('\n'):
                if line:
                    parts = [p.strip() for p in line.split(',')]
                    if len(parts) >= 6:
                        gpus.append({
                            "index": int(parts[0]),
                            "name": parts[1],
                            "memory_total_mb": float(parts[2]),
                            "memory_free_mb": float(parts[3]),
                            "driver_version": parts[4],
                            "compute_capability": parts[5],
                            "type": "NVIDIA",
                            "backend": "CUDA"
                        })
    except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
        pass

    return gpus


def detect_amd_gpus() -> List[Dict[str, Any]]:
    """Detect AMD GPUs using rocm-smi."""
    gpus = []

    try:
        # Try to run rocm-smi
        result = subprocess.run(
            ["rocm-smi", "--showid", "--showmeminfo", "vram"],
            capture_output=True,
            text=True,
            timeout=5
        )

        if result.returncode == 0:
            # Parse rocm-smi output (basic parsing, may need refinement)
            lines = result.st