"""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
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("--force", "-f", is_flag=True, help="Force reinstall even if already installed")
@click.option(
    "--github", "-g", is_flag=True, help="Download from GitHub Releases (faster in China)"
)
def install(backend: str, force: bool, github: bool) -> None:
    """Install PyTorch for a specific backend.

    This command installs 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.1, recommended)
      cuda11   - NVIDIA GPU (CUDA 11.8, for older drivers)
      cuda12   - NVIDIA GPU (CUDA 12.1, 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
    国内用户加速下载 (推荐 --github):
    sagellm install cuda --github      # 从 GitHub 下载 (~800MB，快速)
    sagellm install cuda               # 从官方源下载 (默认，较慢)

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

    console = get_console()
    config = PYTORCH_INSTALL_CONFIGS[backend]

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

    # Show GitHub acceleration hint for CUDA/ROCm backends
    if not github and backend in ("cuda", "cuda11", "cuda12", "haiguang"):
        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]")

        if not force:
            # Check if already working
            ok, msg = config["check"]()
            if ok:
                console.print(f"\n[green]{msg}[/green]")
                console.print("[dim]Already installed! Use -f to force reinstall.[/dim]\n")
                return
    except ImportError:
        console.print("[dim]Current PyTorch: not installed[/dim]")

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

    # Use GitHub Releases if --github flag is set
    github_url = None
    if github and backend in ("cuda", "cuda11", "cuda12", "haiguang"):
        # Map backend to CUDA version (haiguang uses ROCm but wheels are compatible)
        cuda_map = {
            "cuda": "cu121",
            "cuda11": "cu118",
            "cuda12": "cu121",
            "haiguang": "cu121",  # DCU compatible with CUDA wheels
        }
        cuda_ver = cuda_map.get(backend, "cu121")
        pytorch_ver = "2.5.1"  # Current stable version
        github_url = f"https://github.com/intellistream/sagellm-pytorch-wheels/releases/download/v{pytorch_ver}-{cuda_ver}/"
        console.print("\n[blue]🚀 Using GitHub accelerated source:[/blue]")
        console.print(f"[dim]   {github_url}[/dim]")
        console.print("[green]✓ 国内下载速度更快！[/green]\n")

    # Resolve mirror URL (always use official PyPI)
    index_url = None

    # Auto-detect if we need force reinstall (switching between backends)
    need_force = force
    if not need_force:
        try:
            import torch

            current_version = torch.__version__
            target_marker = config.get("version_marker")

            # Determine current backend type from version string
            if "+cpu" in current_version:
                current_type = "cpu"
            elif "+cu" in current_version:
                current_type = "cuda"  # cu118, cu121, etc.
            else:
                current_type = "other"  # Could be Ascend or other

            # Need force reinstall when switching between different backend types
            needs_reinstall = False
            reason = ""

            if backend in ("cuda", "cuda11", "cuda12", "haiguang"):
                # Want CUDA/ROCm, check if current is not CUDA or different version
                if current_type != "cuda":
                    needs_reinstall = True
                    reason = f"切换后端: {current_type} → {config['name']}"
                elif target_marker and target_marker not in current_version:
                    needs_reinstall = True
                    reason = f"切换版本: {current_version} → {target_marker}"
            elif backend == "cpu":
                # Want CPU, check if current is not CPU
                if current_type != "cpu":
                    needs_reinstall = True
                    reason = f"切换后端: {current_type} → CPU"
            elif backend in ("ascend", "kunlun", "muxi"):
                # 国产芯片：检查对应的扩展包是否可用
                ok, _ = config["check"]()
                if not ok:
                    needs_reinstall = True
                    reason = f"{config['name']} 未检测到"

            if needs_reinstall:
                need_force = True
                console.print(f"[yellow]检测到需要切换: {reason}，将强制重新安装[/yellow]")
        except ImportError:
            pass

    # 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]

    if need_force:
        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)
