"""设备检测模块 - 检测本机硬件信息"""
import platform
import subprocess
import re


def detect_device() -> dict:
    """检测本机硬件，返回设备信息字典"""
    info = {
        "device_type": "pc",
        "gpu_model": "",
        "gpu_vram_mb": 0,
        "gpu_compute_units": 0,
        "cpu_model": "",
        "cpu_cores": 0,
        "ram_mb": 0,
        "disk_total_mb": 0,
        "disk_free_mb": 0,
        "disk_type": "unknown",
        "os_info": "",
        "network_type": "unknown",
    }

    # 操作系统
    info["os_info"] = f"{platform.system()} {platform.release()}"

    # CPU
    try:
        import psutil
        info["cpu_cores"] = psutil.cpu_count(logical=True) or os_cpu_count()
        info["cpu_model"] = get_cpu_model()
        info["ram_mb"] = int(psutil.virtual_memory().total / (1024 * 1024))
        # 硬盘
        du = psutil.disk_usage("/")
        info["disk_total_mb"] = int(du.total / (1024 * 1024))
        info["disk_free_mb"] = int(du.free / (1024 * 1024))
        # 尝试检测磁盘类型
        info["disk_type"] = detect_disk_type()
    except ImportError:
        info["cpu_cores"] = os_cpu_count()
        info["cpu_model"] = get_cpu_model()
        info["ram_mb"] = 0

    # GPU
    gpu_info = detect_gpu()
    info["gpu_model"] = gpu_info.get("model", "")
    info["gpu_vram_mb"] = gpu_info.get("vram_mb", 0)
    info["gpu_compute_units"] = gpu_info.get("compute_units", 0)

    return info


def os_cpu_count() -> int:
    """获CPU核心数"""
    try:
        import os
        return os.cpu_count() or 4
    except Exception:
        return 4


def get_cpu_model() -> str:
    """获CPU型号"""
    if platform.system() == "Linux":
        try:
            with open("/proc/cpuinfo") as f:
                for line in f:
                    if "model name" in line:
                        return line.split(":")[1].strip()
        except Exception:
            pass
    elif platform.system() == "Darwin":
        try:
            result = subprocess.run(
                ["sysctl", "-n", "machdep.cpu.brand_string"],
                capture_output=True, text=True, timeout=2
            )
            if result.returncode == 0:
                return result.stdout.strip()
        except Exception:
            pass
    elif platform.system() == "Windows":
        try:
            result = subprocess.run(
                ["wmic", "cpu", "get", "name"],
                capture_output=True, text=True, timeout=2
            )
            lines = result.stdout.strip().split("\n")
            if len(lines) > 1:
                return lines[1].strip()
        except Exception:
            pass
    return platform.processor() or "Unknown"


def detect_gpu() -> dict:
    """检测GPU信息"""
    result = {"model": "", "vram_mb": 0, "compute_units": 0}

    # NVIDIA
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,compute_cap",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5
        )
        if r.returncode == 0 and r.stdout.strip():
            parts = r.stdout.strip().split("\n")[0].split(",")
            result["model"] = parts[0].strip() if len(parts) > 0 else ""
            try:
                result["vram_mb"] = int(float(parts[1].strip())) if len(parts) > 1 else 0
            except ValueError:
                pass
            return result
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # AMD ROCm
    try:
        r = subprocess.run(
            ["rocm-smi", "--showproductname", "--showmeminfo", "vram"],
            capture_output=True, text=True, timeout=5
        )
        if r.returncode == 0:
            for line in r.stdout.split("\n"):
                if "GPU" in line and ":" in line:
                    result["model"] = line.strip()
                if "VRAM Total" in line:
                    m = re.search(r'(\d+)\s*MB', line)
                    if m:
                        result["vram_mb"] = int(m.group(1))
            if result["model"]:
                return result
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # macOS Apple Silicon / Intel
    if platform.system() == "Darwin":
        try:
            r = subprocess.run(
                ["system_profiler", "SPDisplaysDataType"],
                capture_output=True, text=True, timeout=10
            )
            for line in r.stdout.split("\n"):
                if "Chipset Model" in line:
                    result["model"] = line.split(":")[1].strip()
                if "VRAM" in line or "vram" in line:
                    m = re.search(r'(\d+)\s*(MB|GB)', line, re.IGNORECASE)
                    if m:
                        val = int(m.group(1))
                        if m.group(2).upper() == "GB":
                            val *= 1024
                        result["vram_mb"] = val
            # Apple Silicon
            try:
                r2 = subprocess.run(
                    ["sysctl", "-n", "hw.model"],
                    capture_output=True, text=True, timeout=2
                )
                model = r2.stdout.strip()
                if model.startswith("Mac") and not result["model"]:
                    result["model"] = f"Apple {model}"
            except Exception:
                pass
            return result
        except Exception:
            pass

    # Linux - lspci fallback
    if platform.system() == "Linux":
        try:
            r = subprocess.run(
                ["lspci", "-vnn"],
                capture_output=True, text=True, timeout=5
            )
            for line in r.stdout.split("\n"):
                if "VGA" in line or "3D" in line or "Display" in line:
                    if ":" in line:
                        result["model"] = line.split(":")[-1].strip()
                        break
            return result
        except Exception:
            pass

    return result


def detect_disk_type() -> str:
    """检测磁盘类型 (ssd/nvme/hdd)"""
    if platform.system() == "Linux":
        try:
            # 检查nvme
            r = subprocess.run(
                ["lsblk", "-d", "-o", "name,tran,rota"],
                capture_output=True, text=True, timeout=3
            )
            for line in r.stdout.split("\n"):
                parts = line.split()
                if len(parts) >= 2 and parts[0] != "NAME":
                    if "nvme" in parts[1]:
                        return "nvme"
                    if len(parts) >= 3:
                        try:
                            if int(parts[2]) == 0:
                                return "ssd"
                            else:
                                return "hdd"
                        except ValueError:
                            pass
        except Exception:
            pass
    elif platform.system() == "Darwin":
        try:
            r = subprocess.run(
                ["diskutil", "info", "/"],
                capture_output=True, text=True, timeout=5
            )
            if "Solid State" in r.stdout:
                return "ssd"
        except Exception:
            pass
    return "unknown"


if __name__ == "__main__":
    import json
    print(json.dumps(detect_device(), indent=2, ensure_ascii=False))
