"""Packaged CUDA repair helper for sageLLM."""

from __future__ import annotations

import argparse
import os
import shutil
import sys

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

NVIDIA_RUNTIME_PACKAGES = [
    "torch",
    "torchvision",
    "torchaudio",
    "triton",
    "nvidia-nccl-cu12",
    "nvidia-cublas-cu12",
    "nvidia-cuda-runtime-cu12",
    "nvidia-cuda-nvrtc-cu12",
    "nvidia-cudnn-cu12",
    "nvidia-cufft-cu12",
    "nvidia-curand-cu12",
    "nvidia-cusolver-cu12",
    "nvidia-cusparse-cu12",
    "nvidia-cusparselt-cu12",
    "nvidia-cufile-cu12",
    "nvidia-nvjitlink-cu12",
    "nvidia-nvtx-cu12",
    "nvidia-cuda-cupti-cu12",
    "nvidia-nvshmem-cu12",
]

FLASH_ATTN_PACKAGE = "flash-attn>=2.5.0"


def _tail_lines(text: str, limit: int = 20) -> list[str]:
    return [line for line in text.splitlines() if line][-limit:]


def install_flash_attn(
    python_executable: str,
    *,
    allow_source_build: bool,
) -> int:
    print("[3/4] 安装 FlashAttention（NVIDIA 默认增强算子）")
    extra_args = ["--only-binary=:all:"]
    if allow_source_build:
        extra_args = ["--no-build-isolation"]

    result = pip_install_probe(
        python_executable,
        [FLASH_ATTN_PACKAGE],
        extra_args=extra_args,
    )
    if not result.ok:
        print("  ❌ FlashAttention 安装失败")
        if allow_source_build:
            print("  已显式允许源码构建，但安装仍未成功")
        else:
            print("  当前策略禁止默认源码编译 flash-attn，仅允许已验证 binary wheel")
            print("  如需显式允许源码构建，请设置 SAGELLM_ALLOW_FLASH_ATTN_SOURCE_BUILD=1")
        for line in _tail_lines(result.output):
            print(f"    {line}")
        return result.returncode or 1

    verify = run_python_probe(
        python_executable,
        "import flash_attn; print(getattr(flash_attn, '__version__', 'unknown'))",
    )
    if not verify.ok:
        print("  ❌ FlashAttention 导入验证失败")
        for line in _tail_lines(verify.output):
            print(f"    {line}")
        return verify.returncode or 1

    print(f"  ✅ FlashAttention 安装成功: {verify.output}")
    return 0


def diagnose_cuda_runtime(python_executable: str) -> tuple[bool, str, int]:
    script = (
        "from __future__ import annotations\n"
        "try:\n"
        "    import torch\n"
        "except Exception as exc:\n"
        "    print(f'IMPORT_ERROR: {exc.__class__.__name__}: {exc}')\n"
        "    raise SystemExit(1)\n"
        "if not hasattr(torch, 'cuda') or not torch.cuda.is_available() or torch.cuda.device_count() <= 0:\n"
        "    print('CUDA_NOT_READY')\n"
        "    raise SystemExit(2)\n"
        "device_name = torch.cuda.get_device_name(0)\n"
        "torch.zeros(1, device='cuda')\n"
        "print(f'CUDA_READY: {device_name}')\n"
    )
    result = run_python_probe(python_executable, script)
    return result.ok, result.output, result.returncode


def repair_cuda_runtime(
    python_executable: str,
    torch_index_url: str,
    *,
    allow_flash_attn_source_build: bool = False,
) -> int:
    print_banner("CUDA 一键修复与验证")

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

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

    print("[2/3] 重装 PyTorch CUDA 运行时")
    print(f"  使用 PyTorch 索引: {torch_index_url}")
    pip_uninstall_best_effort(python_executable, NVIDIA_RUNTIME_PACKAGES)
    pip_install(
        python_executable,
        ["torch", "torchvision", "torchaudio"],
        extra_args=["--index-url", torch_index_url],
    )
    pip_install(python_executable, ["transformers", "accelerate"])
    print("  ✅ PyTorch CUDA 运行时重装完成")

    flash_attn_exit = install_flash_attn(
        python_executable,
        allow_source_build=allow_flash_attn_source_build,
    )
    if flash_attn_exit != 0:
        return flash_attn_exit

    print("[4/4] 验证 CUDA 导入与设备可用性")
    ok, output, code = diagnose_cuda_runtime(python_executable)
    if output:
        print(output)
    if not ok:
        print("  ❌ CUDA 一键修复后仍未通过验证")
        print("  请重点检查 torch / NCCL / 驱动兼容性，或切换 PyTorch CUDA wheel 索引后重试")
        return code or 1

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


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Repair sageLLM CUDA runtime in-place")
    parser.add_argument(
        "--torch-index-url",
        default=os.environ.get(
            "SAGELLM_CUDA_TORCH_INDEX_URL", "https://download.pytorch.org/whl/cu128"
        ),
        help="PyTorch wheel index to use for reinstall",
    )
    parser.add_argument(
        "--allow-flash-attn-source-build",
        action="store_true",
        default=os.environ.get("SAGELLM_ALLOW_FLASH_ATTN_SOURCE_BUILD") == "1",
        help="Allow source builds for flash-attn when no verified binary wheel is available",
    )
    args = parser.parse_args(argv)
    return repair_cuda_runtime(
        python_executable=sys.executable,
        torch_index_url=args.torch_index_url,
        allow_flash_attn_source_build=args.allow_flash_attn_source_build,
    )


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