"""Packaged Ascend repair helper for sageLLM."""

from __future__ import annotations

import argparse
import shutil
import sys

from sagellm.repair_common import (
    pip_install,
    print_banner,
    refresh_profile,
    run_python_probe,
    run_sagellm_cli,
)

ASCEND_RUNTIME_PACKAGES = ["numpy<2", "decorator", "scipy", "torch", "torch_npu"]


def diagnose_ascend_runtime(python_executable: str) -> tuple[bool, str, int]:
    script = (
        "from __future__ import annotations\n"
        "import importlib\n"
        "try:\n"
        "    import torch\n"
        "    import torch_npu\n"
        "except Exception as exc:\n"
        "    print(f'IMPORT_ERROR: {exc.__class__.__name__}: {exc}')\n"
        "    raise SystemExit(1)\n"
        "try:\n"
        "    importlib.import_module('tbe')\n"
        "except Exception as exc:\n"
        "    print(f'TBE_ERROR: {exc.__class__.__name__}: {exc}')\n"
        "    raise SystemExit(2)\n"
        "if not hasattr(torch, 'npu') or not torch.npu.is_available() or torch.npu.device_count() <= 0:\n"
        "    print('ASCEND_NOT_READY')\n"
        "    raise SystemExit(3)\n"
        "device_name = torch.npu.get_device_name(0)\n"
        "torch.npu.set_device('npu:0')\n"
        "x = torch.ones(1, device='npu')\n"
        "print(f'ASCEND_READY: {device_name} -> {(x + 1).cpu().tolist()}')\n"
    )
    result = run_python_probe(python_executable, script)
    return result.ok, result.output, result.returncode


def repair_ascend_runtime(python_executable: str) -> int:
    print_banner("Ascend 一键修复与验证")

    if shutil.which("npu-smi") is None:
        print("  ❌ 未检测到可用 Ascend NPU（npu-smi 不可用）")
        return 1

    print("[1/4] 诊断当前 Ascend 运行时")
    ok, output, code = diagnose_ascend_runtime(python_executable)
    if output:
        print(output)
    if ok:
        print("  ✅ 当前 Ascend 运行时已可用，但仍将按标准修复链重装以对齐环境")

    print("[2/4] 安装/修复 Ascend Python 运行时")
    pip_install(python_executable, ASCEND_RUNTIME_PACKAGES)
    print("  ✅ Ascend Python 依赖安装完成")

    print("[3/4] 规范化 Ascend 环境与软链")
    fix_cmd = ["ascend", "fix"]
    if shutil.which("bash") is not None:
        fix_cmd.append("--oneclick")
    fix_result = run_sagellm_cli(python_executable, fix_cmd, check=False)
    if fix_result.returncode != 0:
        print("  ⚠️  Ascend fix 命令未完全通过，继续执行最终验证")

    print("[4/4] 验证 Ascend 导入与设备可用性")
    ok, output, code = diagnose_ascend_runtime(python_executable)
    if output:
        print(output)
    if not ok:
        print("  ❌ Ascend 一键修复后仍未通过验证")
        print("  请重点检查 CANN / torch_npu / tbe / LD_LIBRARY_PATH 兼容性")
        return code or 1

    print("  ✅ Ascend 运行时检查通过")
    refresh_profile(python_executable, timeout=60)
    print("  ✅ Ascend 一键修复流程完成")
    return 0


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Repair sageLLM Ascend runtime in-place")
    parser.parse_args(argv)
    return repair_ascend_runtime(python_executable=sys.executable)


if __name__ == "__main__":
    raise SystemExit(main())
