"""Install command - PyTorch backend installation."""

from __future__ import annotations

import logging
import subprocess
import sys

import click

from ..hardware.install_configs import (
    PYTORCH_INSTALL_CONFIGS,
    resolve_backend_index_url,
    resolve_github_release_url,
    run_packaged_backend_repair,
)
from ..utils.console import get_console

logger = logging.getLogger(__name__)


@click.command(hidden=True)
@click.argument(
    "backend",
    type=click.Choice(["cuda", "cuda11", "cuda12", "ascend", "kunlun", "haiguang", "muxi", "cpu"]),
)
@click.option(
    "--github",
    "-g",
    is_flag=True,
    help="Use GitHub Releases when a verified wheel set is published; otherwise fall back to the official index",
)
def install(backend: str, github: bool) -> None:
    """Install PyTorch for a specific backend.

    This command installs or reconciles the correct PyTorch version for your hardware.
    Like Ollama, just run the install command and you're ready to go!

    \b
    国产芯片 (Domestic):
      ascend   - 华为昇腾 Ascend NPU (priority: 100)
      kunlun   - 百度昆仑 Kunlun XPU (priority: 90)
      haiguang - 海光 DCU (priority: 85)
      muxi - 沐曦 Muxi (priority: 80)

    \b
    NVIDIA:
            cuda     - NVIDIA GPU (CUDA 12.8, recommended)
      cuda11   - NVIDIA GPU (CUDA 11.8, for older drivers)
            cuda12   - NVIDIA GPU (CUDA 12.8, same as 'cuda')

    \b
    CPU:
      cpu      - CPU only (smallest download)

    \b
    Examples:
    sagellm install ascend      # 安装华为昇腾支持
    sagellm install kunlun      # 安装百度昆仑支持
    sagellm install haiguang    # 安装海光 DCU 支持
    sagellm install muxi        # 安装沐曦支持
    sagellm install cuda        # 安装 NVIDIA CUDA 支持
    sagellm install cpu         # 安装 CPU-only PyTorch

    \b
    国内用户加速下载 (已发布 wheel 时可用):
    sagellm install cuda --github      # 若对应 wheel 已发布，则走 GitHub 加速
    sagellm install cuda               # 官方 PyTorch 源（默认）

    \b
    GitHub 加速源:
      • 速度更快（GitHub CDN）
      • 官方 wheels，100% 可信
      • 仓库: https://github.com/intellistream/sagellm-pytorch-wheels

        
        CUDA 增强算子策略:
            • NVIDIA 默认额外安装 flash-attn
            • 默认仅接受 binary wheel，不默认源码编译
            • 若确需源码构建，显式设置 SAGELLM_ALLOW_FLASH_ATTN_SOURCE_BUILD=1
    """

    console = get_console()
    config = PYTORCH_INSTALL_CONFIGS[backend]
    github_url = (
        resolve_github_release_url(backend)
        if backend in ("cuda", "cuda11", "cuda12", "haiguang")
        else None
    )

    console.print()
    console.print(f"[bold blue]📦 Installing PyTorch for {config['name']}[/bold blue]")
    console.print()

    if not github:
        repair_exit_code = run_packaged_backend_repair(
            backend,
            python_executable=sys.executable,
        )
        if repair_exit_code is not None:
            raise SystemExit(repair_exit_code)

    # Show GitHub acceleration hint for CUDA/ROCm backends
    if not github and github_url:
        console.print("[yellow]💡 提示：从官方源下载较慢（~800MB）。使用 --github 加速：[/yellow]")
        console.print(f"[dim]   sagellm install {backend} --github[/dim]")
        console.print()

    # Check current status
    try:
        import torch

        current_version = torch.__version__
        console.print(f"[dim]Current PyTorch: {current_version}[/dim]")
    except ImportError:
        console.print("[dim]Current PyTorch: not installed[/dim]")

    # Build pip command
    packages = config["packages"]

    # Use GitHub Releases if --github flag is set
    if github and backend in ("cuda", "cuda11", "cuda12", "haiguang"):
        if github_url:
            console.print("\n[blue]🚀 Using GitHub accelerated source:[/blue]")
            console.print(f"[dim]   {github_url}[/dim]")
            console.print("[green]✓ 国内下载速度更快！[/green]\n")
        else:
            console.print(
                "[yellow]⚠️ 当前后端暂无已验证的 GitHub wheel 加速包，回退官方 PyTorch 源。[/yellow]"
            )
            console.print()

    # Resolve mirror URL (always use official PyPI / PyTorch indices)
    index_url = resolve_backend_index_url(config.get("index_url"))

    # Handle special installation for domestic chips (国产芯片)
    DOMESTIC_BACKENDS = {"ascend", "kunlun", "muxi"}

    if backend in DOMESTIC_BACKENDS:
        setup_guide = config.get("setup_guide", "")
        console.print(f"\n[yellow]⚠️  {config['name']} 需要额外配置：[/yellow]")
        console.print("   1. 安装厂商 SDK/Toolkit")
        console.print("   2. 配置环境变量")
        console.print(f"   3. pip install {' '.join(config['packages'])}")
        if setup_guide:
            console.print(f"\n[dim]详见: {setup_guide}[/dim]\n")

        if not click.confirm(f"是否继续安装 {' 和 '.join(config['packages'])}？"):
            return

        cmd = [sys.executable, "-m", "pip", "install", *config["packages"]]
    elif github_url:
        # Use GitHub Releases with --find-links
        cmd = [
            sys.executable,
            "-m",
            "pip",
            "install",
            *packages,
            "--find-links",
            github_url,
            "--trusted-host",
            "github.com",
        ]
    elif index_url:
        cmd = [
            sys.executable,
            "-m",
            "pip",
            "install",
            *packages,
            "--index-url",
            index_url,
        ]
    else:
        cmd = [sys.executable, "-m", "pip", "install", *packages]

    cmd.append("--force-reinstall")

    # Show command
    cmd_str = " ".join(cmd)
    console.print(f"\n[dim]Running: {cmd_str}[/dim]\n")

    # Show hint for GitHub acceleration if using official source for CUDA
    if not github and backend in ("cuda", "cuda11", "cuda12") and index_url:
        console.print("[yellow]💡 提示：如果下载速度慢，可使用 GitHub 加速源：[/yellow]")
        console.print(f"[dim]   sagellm install {backend} --github[/dim]\n")

    # Run installation
    try:
        subprocess.run(cmd, check=True)
        console.print("\n[green]✅ Installation complete![/green]")

        # Verify installation using a subprocess (can't reload torch in same process)
        console.print("\n[dim]Verifying installation...[/dim]")

        verify_script = config.get("verify_script", "import torch; print(torch.__version__)")
        verify_result = subprocess.run(
            [sys.executable, "-c", verify_script],
            capture_output=True,
            text=True,
        )

        if verify_result.returncode == 0:
            new_version = verify_result.stdout.strip()
            console.print(f"\n[green]✅ PyTorch {new_version} 安装成功！[/green]")

            # Hardware profile detection and caching (Task C)
            console.print("\n[cyan]🔍 Detecting hardware capabilities...[/cyan]")
            try:
                from sagellm_backend.hardware import ProfileCache
                from sagellm_backend.hardware.inspectors import get_inspector

                # Map CLI backend names to hardware types
                backend_to_hw_type = {
                    "cuda": "cuda",
                    "cuda11": "cuda",
                    "cuda12": "cuda",
                    "ascend": "ascend",
                    "cpu": "cpu",
                    "haiguang": "cuda",  # DCU compatible with CUDA
                    # Add other backends as needed
                }
                hw_type = backend_to_hw_type.get(backend, backend)

                inspector = get_inspector(hw_type)
                # Include functional tests for full verification (L3)
                profile = inspector.get_profile(include_functional=True)

                cache = ProfileCache()
                cache.update(hw_type, profile)

                console.print(f"[green]✅ Hardware profile saved to {cache.cache_path}[/green]")

                # Display detection results
                if profile.is_usable():
                    console.print(f"[bold green]✅ {hw_type.upper()} is ready![/bold green]")
                    console.print(f"   Max Level: {profile.max_level.value}")
                    console.print(f"   Devices: {profile.device_count}")
                    if profile.device_names:
                        console.print(f"   Device: {profile.device_names[0]}")
                else:
                    console.print(f"[yellow]⚠️  {hw_type.upper()} partially available[/yellow]")
                    console.print(f"   Max Level: {profile.max_level.value}")
                    if profile.l2_framework.suggestions:
                        console.print("[cyan]Suggestions:[/cyan]")
                        for s in profile.l2_framework.suggestions[:3]:  # Show first 3
                            console.print(f"  - {s}")
            except Exception as e:
                # Don't fail installation if detection fails
                logger.debug("Hardware detection failed: %s", e)
                console.print(f"[dim]Hardware detection skipped: {e}[/dim]")

            console.print("\n[bold green]🎉 Ready to use![/bold green]")
            console.print(f"[dim]Try: sagellm run -p 'Hello' --backend={backend}[/dim]\n")
        else:
            console.print("\n[yellow]⚠️  Installation complete but verification failed.[/yellow]")
            console.print(f"[dim]{verify_result.stderr}[/dim]")
            console.print("[dim]Please restart your terminal and try again.[/dim]\n")

    except subprocess.CalledProcessError as e:
        console.print(f"\n[red]❌ Installation failed: {e}[/red]")
        console.print("[dim]Please check your network and try again.[/dim]\n")
        sys.exit(1)
