from __future__ import annotations

import multiprocessing
import platform
import subprocess

from epochler.contract.models import HardwareProfile


def detect_hardware() -> HardwareProfile:
    """Detect available hardware: CPU, RAM, GPU (CUDA/MPS)."""
    profile = HardwareProfile()

    profile.cpu_count = multiprocessing.cpu_count()

    try:
        import psutil
        profile.ram_total_gb = round(psutil.virtual_memory().total / (1024 ** 3), 1)
    except ImportError:
        if platform.system() == "Darwin":
            try:
                result = subprocess.run(
                    ["sysctl", "-n", "hw.memsize"],
                    capture_output=True, text=True, timeout=5,
                )
                profile.ram_total_gb = round(int(result.stdout.strip()) / (1024 ** 3), 1)
            except Exception:
                pass
        else:
            try:
                with open("/proc/meminfo") as f:
                    for line in f:
                        if line.startswith("MemTotal"):
                            kb = int(line.split()[1])
                            profile.ram_total_gb = round(kb / (1024 ** 2), 1)
                            break
            except Exception:
                pass

    # CUDA detection
    try:
        import torch
        if torch.cuda.is_available():
            profile.gpu_available = True
            profile.gpu_type = "cuda"
            profile.gpu_count = torch.cuda.device_count()
            for i in range(profile.gpu_count):
                profile.gpu_names.append(torch.cuda.get_device_name(i))
                vram = torch.cuda.get_device_properties(i).total_mem / (1024 ** 3)
                profile.gpu_vram_gb.append(round(vram, 1))
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            profile.gpu_available = True
            profile.gpu_type = "mps"
            profile.gpu_count = 1
            profile.gpu_names = ["Apple Silicon (MPS)"]
            profile.gpu_vram_gb = [profile.ram_total_gb]
    except ImportError:
        # Try nvidia-smi as fallback for CUDA detection without torch
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
                capture_output=True, text=True, timeout=10,
            )
            if result.returncode == 0:
                profile.gpu_available = True
                profile.gpu_type = "cuda"
                lines = result.stdout.strip().splitlines()
                profile.gpu_count = len(lines)
                for line in lines:
                    parts = [p.strip() for p in line.split(",")]
                    profile.gpu_names.append(parts[0])
                    profile.gpu_vram_gb.append(round(float(parts[1]) / 1024, 1))
        except Exception:
            pass

    return profile
