# -*- coding: utf-8 -*-
"""华为昇腾 NPU 推理后端实现（ACL/CANN）。

对应 PRD 5.6 与 issue #59「昇腾 NPU 后端适配（CANN 对接）」：
昇腾实现（ACL/CANN）通过昇腾推理服务（MindIE / onnxruntime-ascend）的
OpenAI 兼容接口对外提供推理，与 NVIDIA GPU 后端实现同一
`InferenceBackend` 接口：**切换后端仅改适配层配置，业务代码零改动**。

两种运行模式（由 endpoint 是否配置决定）：
1. **服务模式**（endpoint 非空）：对接 MindIE / onnxruntime-ascend 的
   OpenAI 兼容 HTTP 接口，走 ``/acl/models/load``、``/v1/chat/completions``、
   ``/acl/models/unload``、``/health``；
2. **直连模式**（endpoint 为空且本机可检测到 CANN）：调用本地 ACL Python
   API（``acl.init`` / ``acl.rt.set_device`` / ``acl.rt.reset_device`` /
   ``acl.finalize``）管理昇腾设备与模型加载；
   endpoint 为空且无 CANN 环境时进入 dry-run（适配层就绪，便于离线验证）。

健康巡检会附带 :func:`~inference_backend.cann_probe.probe_cann_environment`
探测到的 CANN 设备/工具链信息（``cann`` 段），便于运维确认后端就绪状态。
"""
import time

from inference_backend.base import InferenceBackend, InferRequest, InferResult
from inference_backend.cann_probe import probe_cann_environment


class AscendNpuBackend(InferenceBackend):
    """华为昇腾 NPU 后端：面向昇腾 310P/910B（CANN/MindIE）。"""

    backend_name = "npu"

    def __init__(self, endpoint: str = "", model: str = "iaop-default",
                 timeout_seconds: float = 10.0, runtime: str = "mindie",
                 device: str = "ascend-910b", cann_version: str = "8.0",
                 **kwargs):
        super().__init__(endpoint, model, timeout_seconds)
        self.runtime = runtime            # mindie | onnx-ascend
        self.device = device              # ascend-310p | ascend-910b
        self.cann_version = cann_version  # CANN 工具链版本
        self._loaded = False
        self._acl_initialized = False
        self._cann_cache = None

    # ---- CANN 环境探测（懒加载，探测结果缓存于实例） ----

    def cann_info(self) -> dict:
        """返回本机 CANN 环境探测结果（见 cann_probe.probe_cann_environment）。"""
        if self._cann_cache is None:
            self._cann_cache = probe_cann_environment()
        return self._cann_cache

    # ---- 统一接口实现 ----

    def load_model(self, model_name: str | None = None) -> dict:
        """加载模型到 NPU。

        服务模式：``POST /acl/models/load``（ACL aclmdlLoadFromFile 语义）；
        直连模式：本地 ``acl.init()`` + ``acl.rt.set_device(device_id)``；
        无 endpoint 且无 CANN 环境：dry-run，仅上报适配层就绪。
        """
        model_name = model_name or self.model
        cann = self.cann_info()
        if self.endpoint:
            body = self._post_json(
                "/acl/models/load",
                {"model": model_name, "device": self.device,
                 "cann_version": self.cann_version})
            self._loaded = body.get("status") in ("ok", "loaded", "ready")
            return body
        if cann["available"]:
            self._acl_load()
            self._loaded = True
            return {"status": "ok", "backend": self.backend_name,
                    "model": model_name, "device": self.device,
                    "runtime": self.runtime, "cann": cann,
                    "reason": "直连模式：CANN ACL 已初始化并绑定设备"}
        self._loaded = True
        return {"status": "ok", "backend": self.backend_name,
                "model": model_name, "device": self.device,
                "cann": cann,
                "reason": "dry-run（未配置 endpoint 且无 CANN 环境，适配层就绪）"}

    def _acl_load(self) -> None:
        """直连模式初始化 CANN ACL 并把上下文绑定到首张昇腾设备。"""
        try:
            import acl
        except Exception as exc:  # 探测与实际导入之间环境可能变化，容错
            raise RuntimeError(f"CANN ACL 不可用，无法直连加载: {exc}") from exc
        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):
            raise RuntimeError(f"acl.init() 失败: ret={ret}")
        device_id = 0  # 默认首卡；多卡资源调度由部署侧配置扩展
        ret = acl.rt.set_device(device_id)
        if not (ret is True or ret is None or ret == 0):
            acl.finalize()
            raise RuntimeError(f"acl.rt.set_device({device_id}) 失败: ret={ret}")
        self._acl_initialized = True

    def infer(self, request: InferRequest) -> InferResult:
        """昇腾推理（MindIE OpenAI 兼容 /v1/chat/completions）。"""
        started = time.monotonic()
        payload = request.to_payload()
        payload["model"] = payload["model"] or self.model
        body = self._post_json("/v1/chat/completions", payload)
        latency_ms = round((time.monotonic() - started) * 1000, 2)
        try:
            text = body["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError):
            text = str(body)
        return InferResult(
            text=text,
            backend=self.backend_name,
            latency_ms=latency_ms,
            meta={"runtime": self.runtime, "device": self.device,
                  "cann_version": self.cann_version, "model": self.model,
                  "cann": self.cann_info(), "raw": body},
        )

    def health(self) -> dict:
        """健康巡检：探测 /health，返回后端/设备/CANN 环境信息。"""
        base = self._healthz()
        base.update({
            "backend": self.backend_name,
            "runtime": self.runtime,
            "device": self.device,
            "cann_version": self.cann_version,
            "model": self.model,
            "loaded": self._loaded,
            "cann": self.cann_info(),
        })
        return base

    def unload(self) -> dict:
        """卸载模型、释放 NPU 资源。

        服务模式：``POST /acl/models/unload``（ACL aclmdlUnload 语义）；
        直连模式：``acl.rt.reset_device`` + ``acl.finalize``；
        无 endpoint 且未直连加载：dry-run。
        """
        self._loaded = False
        if self.endpoint:
            return self._post_json("/acl/models/unload", {"model": self.model})
        if self._acl_initialized:
            try:
                import acl
                acl.rt.reset_device(0)
                acl.finalize()
            except Exception as exc:  # noqa: BLE001
                return {"status": "warn", "backend": self.backend_name,
                        "reason": f"ACL 资源释放失败: {exc}"}
            self._acl_initialized = False
            return {"status": "ok", "backend": self.backend_name,
                    "reason": "直连模式：ACL 资源已释放"}
        return {"status": "ok", "backend": self.backend_name,
                "reason": "dry-run（未配置 endpoint）"}