"""Auto-discovery of hardware on a mesh node."""
from __future__ import annotations

import logging
import multiprocessing
import os
import platform
import re
import shutil
import subprocess
from pathlib import Path

import psutil

from mesh.models import CPUInfo, FPGAInfo, GPUInfo, HardwareInventory, MountInfo

logger = logging.getLogger(__name__)

_SLENTO_REQUIRED_LIBS = [
    "pytorch_compat/libcudart.so.12",
    "pytorch_compat/libcublas.so.12",
    "pytorch_compat/libcublasLt.so.12",
    "pytorch_compat/libcudnn.so.9",
    "pytorch_compat/libcuda.so.1",
]
_SLENTO_GPU_ARCH_MAP = {
    "7900 XTX": "gfx1100",
    "7900 XT": "gfx1100",
    "MI100": "gfx908",
    "MI250": "gfx90a",
    "MI300": "gfx942",
}


def scan_hardware(hostname_override: str = "", config=None) -> HardwareInventory:
    """Detect all hardware on this machine."""
    inv = HardwareInventory(
        hostname=hostname_override or platform.node(),
        platform=platform.system(),
        memory_total_mb=psutil.virtual_memory().total // (1024 * 1024),
    )
    inv.cpu = _detect_cpu()
    inv.gpus = _detect_gpus()
    inv.fpgas = _detect_fpgas()
    inv.has_pytorch = _check_pytorch()
    inv.has_rocm = shutil.which("rocm-smi") is not None
    inv.has_cuda = shutil.which("nvidia-smi") is not None
    inv.mounts = _detect_mounts()
    inv.scratch_paths = _get_config_value(config, "scratch_paths", []) or []

    slento = _detect_slento_cuda(config=config)
    inv.has_slento_cuda = slento["has_slento_cuda"]
    inv.slento_cuda_version = slento["slento_cuda_version"]
    inv.slento_cuda_gpu_archs = slento["slento_cuda_gpu_archs"]
    inv.slento_cuda_lib_path = slento["slento_cuda_lib_path"]
    return inv


def _get_config_value(config, key: str, default=None):
    if config is None:
        return default
    if isinstance(config, dict):
        if key in config:
            return config[key]
        node_cfg = config.get("node")
        if isinstance(node_cfg, dict) and key in node_cfg:
            return node_cfg[key]
        return default
    if hasattr(config, key):
        return getattr(config, key)
    node_cfg = getattr(config, "node", None)
    if node_cfg is not None and hasattr(node_cfg, key):
        return getattr(node_cfg, key)
    return default


def _detect_cpu() -> CPUInfo:
    info = CPUInfo(cores=multiprocessing.cpu_count() or 1)
    # Model name
    if platform.system() == "Linux":
        try:
            with open("/proc/cpuinfo") as f:
                for line in f:
                    if "model name" in line:
                        info.model = line.split(":")[1].strip()
                        break
        except Exception:
            pass
    elif platform.system() == "Darwin":
        try:
            info.model = subprocess.check_output(
                ["sysctl", "-n", "machdep.cpu.brand_string"],
                text=True, stderr=subprocess.DEVNULL
            ).strip()
        except Exception:
            pass

    # lscpu for Linux details
    try:
        out = subprocess.check_output("lscpu", text=True, stderr=subprocess.DEVNULL)
        cores_per = 1
        sockets = 1
        for line in out.splitlines():
            if "Core(s) per socket:" in line:
                cores_per = int(line.split(":")[1].strip())
            elif "Socket(s):" in line:
                sockets = int(line.split(":")[1].strip())
            elif "Thread(s) per core:" in line:
                info.threads_per_core = int(line.split(":")[1].strip())
            elif "NUMA node(s):" in line:
                info.numa_nodes = int(line.split(":")[1].strip())
            elif "CPU max MHz:" in line:
                info.freq_mhz = float(line.split(":")[1].strip())
            elif "CPU MHz:" in line and info.freq_mhz == 0:
                info.freq_mhz = float(line.split(":")[1].strip())
        info.physical_cores = cores_per * sockets
    except Exception:
        info.physical_cores = info.cores // 2

    # Cache sizes from sysfs
    try:
        for idx in range(10):
            cache_dir = f"/sys/devices/system/cpu/cpu0/cache/index{idx}"
            if not os.path.exists(cache_dir):
                break
            level = open(f"{cache_dir}/level").read().strip()
            ctype = open(f"{cache_dir}/type").read().strip()
            size_str = open(f"{cache_dir}/size").read().strip()
            size_kb = int(re.match(r"(\d+)", size_str).group(1))
            if "M" in size_str:
                size_kb *= 1024
            if level == "1" and "Data" in ctype:
                info.l1d_kb = size_kb
            elif level == "2":
                info.l2_kb = size_kb
            elif level == "3":
                info.l3_kb = size_kb
    except Exception:
        pass
    return info


def _detect_gpus() -> list[GPUInfo]:
    gpus = []
    # NVIDIA GPUs
    gpus.extend(_detect_nvidia_gpus())
    # AMD GPUs
    gpus.extend(_detect_amd_gpus())
    return gpus


def _detect_mounts() -> list[MountInfo]:
    mounts = []
    try:
        for partition in psutil.disk_partitions(all=False):
            try:
                usage = psutil.disk_usage(partition.mountpoint)
            except Exception:
                continue
            mounts.append(MountInfo(
                mountpoint=partition.mountpoint,
                device=partition.device,
                fstype=partition.fstype,
                total_mb=int(usage.total // (1024 * 1024)),
                free_mb=int(usage.free // (1024 * 1024)),
            ))
    except Exception as e:
        logger.debug("Mount detection failed: %s", e)
    return mounts


def _detect_nvidia_gpus() -> list[GPUInfo]:
    gpus = []
    if not shutil.which("nvidia-smi"):
        return gpus
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=name,memory.total,driver_version,temperature.gpu,utilization.gpu,memory.used",
             "--format=csv,noheader,nounits"],
            text=True, stderr=subprocess.DEVNULL, timeout=10
        )
        for line in out.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 6:
                gpus.append(GPUInfo(
                    name=parts[0],
                    vendor="nvidia",
                    vram_mb=int(float(parts[1])),
                    driver=parts[2],
                    temperature_c=float(parts[3]) if parts[3] != "N/A" else 0.0,
                    utilization_pct=float(parts[4]) if parts[4] != "N/A" else 0.0,
                    memory_used_mb=int(float(parts[5])) if parts[5] != "N/A" else 0,
                ))
    except Exception as e:
        logger.warning("nvidia-smi failed: %s", e)
    return gpus


def _detect_amd_gpus() -> list[GPUInfo]:
    gpus = []
    if not shutil.which("rocm-smi"):
        return gpus
    try:
        out = subprocess.check_output(
            ["rocm-smi", "--showproductname", "--showmeminfo", "vram",
             "--showtemp", "--showuse", "--json"],
            text=True, stderr=subprocess.DEVNULL, timeout=10
        )
        import json
        data = json.loads(out)
        for key, val in data.items():
            if not key.startswith("card"):
                continue
            gpu = GPUInfo(vendor="amd")
            gpu.name = val.get("Card Series", val.get("Card series", "AMD GPU"))
            vram_total = val.get("VRAM Total Memory (B)", 0)
            if vram_total:
                gpu.vram_mb = int(vram_total) // (1024 * 1024)
            vram_used = val.get("VRAM Total Used Memory (B)", 0)
            if vram_used:
                gpu.memory_used_mb = int(vram_used) // (1024 * 1024)
            temp = val.get("Temperature (Sensor edge) (C)", 0)
            if temp:
                gpu.temperature_c = float(temp)
            use = val.get("GPU use (%)", 0)
            if use:
                gpu.utilization_pct = float(use)
            gpus.append(gpu)
    except Exception as e:
        logger.warning("rocm-smi failed: %s", e)
    # Fallback: check /sys/class/drm
    if not gpus:
        try:
            drm = Path("/sys/class/drm")
            for card in sorted(drm.glob("card[0-9]*")):
                device = card / "device"
                vendor_file = device / "vendor"
                if vendor_file.exists():
                    vendor_id = vendor_file.read_text().strip()
                    if vendor_id == "0x1002":  # AMD
                        name = "AMD GPU"
                        product = device / "product_name"
                        if product.exists():
                            name = product.read_text().strip()
                        vram = 0
                        mem_file = device / "mem_info_vram_total"
                        if mem_file.exists():
                            vram = int(mem_file.read_text().strip()) // (1024 * 1024)
                        gpus.append(GPUInfo(name=name, vendor="amd", vram_mb=vram))
        except Exception:
            pass
    return gpus


def _detect_fpgas() -> list[FPGAInfo]:
    fpgas = []
    try:
        out = subprocess.check_output(
            ["lspci", "-nn"], text=True, stderr=subprocess.DEVNULL, timeout=5
        )
        for line in out.splitlines():
            lower = line.lower()
            if "xilinx" in lower or "altera" in lower or "intel.*fpga" in lower:
                bdf = line.split()[0]
                vendor = "xilinx" if "xilinx" in lower else "altera/intel"
                name = line.split(": ", 1)[1] if ": " in line else "FPGA"
                # Check driver
                driver = ""
                driver_link = Path(f"/sys/bus/pci/devices/0000:{bdf}/driver")
                if driver_link.is_symlink():
                    driver = driver_link.resolve().name
                fpgas.append(FPGAInfo(
                    name=name, vendor=vendor, pcie_bdf=bdf, driver=driver
                ))
    except Exception:
        pass
    return fpgas


def _check_pytorch() -> bool:
    try:
        import torch
        return True
    except ImportError:
        return False


def _detect_slento_cuda(config=None) -> dict[str, object]:
    result = {
        "has_slento_cuda": False,
        "slento_cuda_version": "",
        "slento_cuda_gpu_archs": [],
        "slento_cuda_lib_path": "",
    }

    if platform.system() != "Linux":
        return result
    if not shutil.which("rocm-smi"):
        return result

    amd_gpus = _detect_amd_gpus()
    if not amd_gpus:
        return result

    script_path = _find_slento_script(config=config)
    if not script_path:
        return result

    base_dir = script_path.parent
    missing = [rel for rel in _SLENTO_REQUIRED_LIBS if not (base_dir / rel).exists()]
    if missing:
        logger.debug("Slento CUDA candidate missing libs under %s: %s", base_dir, missing)
        return result

    result["has_slento_cuda"] = True
    result["slento_cuda_version"] = _detect_slento_version(base_dir)
    result["slento_cuda_gpu_archs"] = _detect_slento_gpu_archs()
    result["slento_cuda_lib_path"] = str(base_dir / "pytorch_compat")
    return result


def _find_slento_script(config=None) -> Path | None:
    configured = _get_config_value(config, "slento_cuda_path", "")
    candidates: list[Path] = []
    if configured:
        cfg_path = Path(configured).expanduser()
        if cfg_path.name == "cuda_emulator_run.sh":
            candidates.append(cfg_path)
        else:
            candidates.append(cfg_path / "cuda_emulator_run.sh")
    candidates.extend([
        Path("~/CascadeProjects/rocm-cuda-compat/build/cuda_emulator_run.sh").expanduser(),
        Path("/opt/slento-cuda/cuda_emulator_run.sh"),
    ])
    which_path = shutil.which("cuda_emulator_run.sh")
    if which_path:
        candidates.append(Path(which_path))

    for candidate in candidates:
        if candidate.exists() and candidate.is_file():
            return candidate.resolve()
    return None


def _detect_slento_version(base_dir: Path) -> str:
    for version_file in ("VERSION", "version.txt"):
        path = base_dir / version_file
        if path.exists():
            try:
                return path.read_text().strip()
            except Exception:
                pass
    if re.search(r"slento", str(base_dir), re.IGNORECASE):
        return base_dir.name
    return "detected"


def _detect_slento_gpu_archs() -> list[str]:
    archs: list[str] = []
    try:
        out = subprocess.check_output(
            ["rocm-smi", "--showproductname"],
            text=True,
            stderr=subprocess.DEVNULL,
            timeout=10,
        )
        for product_name, arch in _SLENTO_GPU_ARCH_MAP.items():
            if product_name in out and arch not in archs:
                archs.append(arch)
    except Exception as e:
        logger.debug("rocm-smi product lookup failed for Slento CUDA detection: %s", e)
    return archs
