# -*- coding: utf-8 -*-
"""昇腾 CANN 运行时环境探测（issue #59「昇腾 NPU 后端适配（CANN 对接）」）。

目标：
- 在**无 CANN 硬件的开发/CI 环境**下安全降级（available=False，绝不抛异常）；
- 在**有 CANN 的生产环境**正确上报设备数量/名称与工具链版本，
  供 :class:`~inference_backend.npu_backend.AscendNpuBackend` 的健康巡检
  与直连模式（本地 ACL 生命周期管理）使用。

探测优先级（取第一个可用的运行时）：
1. ``torch_npu`` —— PyTorch 昇腾插件（``torch.npu.*``）；
2. ``acl`` —— CANN 基础 ACL Python API（``acl.init`` / ``acl.rt.*``）；
3. ``npu-smi`` —— 昇腾命令行工具（``npu-smi info`` 输出解析）。

ACL API 各 CANN 版本存在差异（部分接口返回 ``(ret, value)`` 元组、
部分直接返回值），本模块统一用 ``_unwrap`` 兼容两种风格，并逐调用
try/except 兜底 —— 探测逻辑在任何版本/异常下都不会让调用方崩溃。
"""
from __future__ import annotations

import shutil
import subprocess
from typing import Optional

#: 探测失败兜底结果（永不抛异常）
_NO_CANN = {
    "available": False,
    "runtime": "none",
    "cann_version": "",
    "device_count": 0,
    "device_names": [],
    "detail": "未检测到 CANN 运行时（torch_npu / acl / npu-smi 均不可用）",
}


def probe_cann_environment() -> dict:
    """探测当前环境的昇腾 CANN 运行时，返回结构化结果。

    Returns:
        dict，固定键：``available`` / ``runtime`` / ``cann_version`` /
        ``device_count`` / ``device_names`` / ``detail``。
        任一探测路径失败均不抛异常，仅将 ``available`` 置为 False。
    """
    for probe in (_probe_torch_npu, _probe_acl, _probe_npu_smi):
        info = probe()
        if info is not None:
            return info
    return dict(_NO_CANN)


def _unwrap(ret_value):
    """兼容 ACL API 两种返回风格：``(ret, value)`` 元组或直接返回 value。"""
    if isinstance(ret_value, tuple):
        if len(ret_value) >= 2:
            return ret_value[1]
        return None
    return ret_value


def _version_of(module, name: str) -> str:
    """尽力取模块版本号（torch_npu.__version__ / acl.__version__）。"""
    try:
        ver = getattr(module, "__version__", "")
        return str(ver) if ver else ""
    except Exception:
        return ""


def _probe_torch_npu() -> Optional[dict]:
    """探测 torch_npu（PyTorch 昇腾插件）路径。"""
    try:
        import torch  # noqa: F401
        import torch_npu  # noqa: F401
    except Exception:
        return None  # 未安装该运行时，交给下一个探测路径
    base = {
        "runtime": "torch_npu",
        "cann_version": _version_of(torch_npu, "torch_npu"),
    }
    try:
        if not torch_npu.npu.is_available():
            return {
                **base,
                "available": False,
                "device_count": 0,
                "device_names": [],
                "detail": "torch_npu 可导入但 NPU 不可用（未检测到昇腾设备）",
            }
        count = int(torch.npu.device_count())
        names = []
        for i in range(count):
            try:
                names.append(str(torch.npu.get_device_name(i)))
            except Exception:
                names.append("ascend-device-%d" % i)
        return {
            **base,
            "available": True,
            "device_count": count,
            "device_names": names,
            "detail": "torch_npu 探测成功（%d 张昇腾设备）" % count,
        }
    except Exception as exc:  # noqa: BLE001 —— 探测必须容错
        return {
            **base,
            "available": False,
            "device_count": 0,
            "device_names": [],
            "detail": "torch_npu 探测失败: %s" % exc,
        }


def _probe_acl() -> Optional[dict]:
    """探测 CANN ACL Python API 路径。"""
    try:
        import acl
    except Exception:
        return None  # 未安装该运行时，交给下一个探测路径
    base = {
        "runtime": "acl",
        "cann_version": _version_of(acl, "acl"),
    }
    try:
        ret = acl.init()
        # 注意：不能用 `ret not in (0, None, True)` —— Python 中 1 == True，
        # 会把错误码 1 误判为成功；ACL 约定 ret=0（ACL_SUCCESS）为成功。
        if not (ret is True or ret is None or ret == 0):
            return {
                **base,
                "available": False,
                "device_count": 0,
                "device_names": [],
                "detail": "acl.init() 返回错误码 %s" % ret,
            }
        try:
            count = int(_unwrap(acl.rt.get_device_count()))
        except Exception:
            count = 0
        names = []
        for i in range(count):
            name = None
            for method in ("get_device_name", "get_soc_name"):
                fn = getattr(acl.rt, method, None)
                if fn is None:
                    continue
                try:
                    name = _unwrap(fn(i))
                    break
                except Exception:
                    name = None
            names.append(str(name) if name else "ascend-device-%d" % i)
        # CANN 版本：优先 acl.rt.get_version()，其次模块 __version__
        try:
            ver = _unwrap(acl.rt.get_version()) or base["cann_version"]
        except Exception:
            ver = base["cann_version"]
        return {
            **base,
            "cann_version": str(ver) if ver else "",
            "available": True,
            "device_count": count,
            "device_names": names,
            "detail": "ACL 探测成功（%d 张昇腾设备）" % count,
        }
    except Exception as exc:  # noqa: BLE001
        return {
            **base,
            "available": False,
            "device_count": 0,
            "device_names": [],
            "detail": "ACL 探测失败: %s" % exc,
        }
    finally:
        try:
            acl.finalize()  # 探测后立即释放，避免占用设备
        except Exception:
            pass


def _probe_npu_smi() -> Optional[dict]:
    """探测 npu-smi 命令行工具路径。"""
    exe = shutil.which("npu-smi")
    if not exe:
        return None  # 未安装该工具
    try:
        proc = subprocess.run(
            [exe, "info", "-l"],
            capture_output=True, text=True, timeout=10,
            check=False,
        )
        text = proc.stdout + "\n" + proc.stderr
    except Exception as exc:
        return {
            "available": False,
            "runtime": "npu-smi",
            "cann_version": "",
            "device_count": 0,
            "device_names": [],
            "detail": "npu-smi 执行失败: %s" % exc,
        }
    # 解析 "Device Count : N" 与 "Name : xxx"（npu-smi info -l 常见输出）
    device_count = 0
    names = []
    for line in text.splitlines():
        low = line.lower()
        if "device count" in low and ":" in line:
            try:
                device_count = int(line.split(":")[-1].strip())
            except ValueError:
                device_count = 0
        elif "name" in low and ":" in line:
            name = line.split(":", 1)[-1].strip()
            if name:
                names.append(name)
    available = device_count > 0 or bool(names)
    return {
        "available": available,
        "runtime": "npu-smi",
        "cann_version": "",
        "device_count": device_count,
        "device_names": names[: max(device_count, len(names))],
        "detail": "npu-smi 探测成功（%d 张昇腾设备）" % device_count
        if available else "npu-smi 可执行但未解析到昇腾设备",
    }