"""Domestic chip hardware detection (Kunlun, Haiguang, Muxi)."""

from __future__ import annotations

import os
import subprocess
import sys


def _run_probe(script: str, timeout_sec: int = 20) -> tuple[bool, str]:
    env = os.environ.copy()
    env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        timeout=timeout_sec,
        env=env,
    )
    if result.returncode != 0:
        reason = (result.stderr or result.stdout or "probe failed").strip()
        return False, reason.splitlines()[-1] if reason else "probe failed"
    return True, (result.stdout or "").strip()


def check_kunlun_available() -> tuple[bool, str]:
    """Check if Kunlun XPU is available."""
    try:
        ok, output = _run_probe(
            "import torch; "
            "import torch_xmlir; "
            "ready = hasattr(torch, 'xpu') and torch.xpu.is_available(); "
            "print('1' if ready else '0')"
        )
        if not ok:
            if "No module named" in output and "torch_xmlir" in output:
                return False, "❌ torch_xmlir not installed"
            return False, f"❌ Error: {output}"
        return (
            (True, "✅ Kunlun XPU ready")
            if output.startswith("1")
            else (False, "❌ Kunlun XPU not detected")
        )
    except subprocess.TimeoutExpired:
        return False, "❌ Error: Kunlun probe timeout"
    except Exception as e:
        return False, f"❌ Error: {e}"


def check_haiguang_available() -> tuple[bool, str]:
    """Check if Haiguang DCU is available."""
    try:
        ok, output = _run_probe(
            "import torch; "
            "ready = torch.cuda.is_available(); "
            "name = torch.cuda.get_device_name(0) if ready else ''; "
            "lname = name.lower(); "
            "hit = ready and ('dcu' in lname or 'hygon' in lname or 'gfx' in lname); "
            "print('1:' + name if hit else '0')"
        )
        if not ok:
            if "No module named" in output and "torch" in output:
                return False, "❌ PyTorch not installed"
            return False, f"❌ Error: {output}"
        if output.startswith("1:"):
            return True, f"✅ Haiguang DCU ready: {output[2:] or 'DCU device'}"
        return False, "❌ Haiguang DCU not detected"
    except subprocess.TimeoutExpired:
        return False, "❌ Error: Haiguang probe timeout"
    except Exception as e:
        return False, f"❌ Error: {e}"


def check_muxi_available() -> tuple[bool, str]:
    """Check if Muxi GPU is available."""
    try:
        ok, output = _run_probe(
            "import torch; "
            "import torch_muxi; "
            "ready = hasattr(torch, 'muxi') and torch.muxi.is_available(); "
            "cnt = torch.muxi.device_count() if ready else 0; "
            "print(f'1:{cnt}' if ready else '0')"
        )
        if not ok:
            if "No module named" in output and "torch_muxi" in output:
                return False, "❌ torch_muxi not installed"
            return False, f"❌ Error: {output}"
        if output.startswith("1:"):
            count = output[2:] or "1"
            return True, f"✅ Muxi GPU ready ({count} device(s))"
        return False, "❌ Muxi GPU not detected"
    except subprocess.TimeoutExpired:
        return False, "❌ Error: Muxi probe timeout"
    except Exception as e:
        return False, f"❌ Error: {e}"
