"""hello command - one-click inference test."""

import os
import subprocess
import sys

import click


@click.command()
@click.option("--cpu", is_flag=True, help="Force CPU backend (skip CUDA detection)")
def hello(cpu: bool) -> None:
    """🚀 One-click inference test - verify installation works!

    Downloads a small model and runs a simple inference.
    Auto-detects CUDA GPU and uses it if available.

    \b
    Example:
      pip install isagellm
    sagellm hello            # Auto-detect (uses GPU if available)
    sagellm hello --cpu      # Force CPU
    """
    from ..utils.backend import detect_ascend
    from ..utils.console import get_console
    from ..utils.env_detector import EnvInitializer

    def _exit_with_code(code: int) -> None:
        if os.environ.get("SAGELLM_HELLO_FORCE_CLEAN_EXIT", "1") == "1":
            try:
                sys.stdout.flush()
                sys.stderr.flush()
            finally:
                os._exit(code)
        sys.exit(code)

    os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"

    console = get_console()
    console.print("\n[bold blue]🚀 sageLLM Hello World[/bold blue]\n")

    # Re-detect HF endpoint before downloading — handles VPN coming online
    # after the profile cache was written (stale mirror cache on VPN machines).
    EnvInitializer.init_hf_endpoint(force_detect=True, verbose=False)

    # Detect backend
    backend_kind = "cpu"
    torch_probe_error: Exception | None = None
    ascend_probe_timeout_sec = int(os.getenv("SAGELLM_ASCEND_PROBE_TIMEOUT_SEC", "25"))

    def _ascend_functional_ready() -> tuple[bool, str | None]:
        probe_script = (
            "import os; "
            "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD']='0'; "
            "import torch, torch_npu; "
            "assert hasattr(torch, 'npu') and torch.npu.is_available(); "
            "torch.npu.set_device('npu:0'); "
            "x=torch.ones(1, device='npu'); "
            "_=(x+1).cpu(); "
            "print('ok')"
        )
        try:
            env = os.environ.copy()
            env["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
            result = subprocess.run(
                [sys.executable, "-c", probe_script],
                capture_output=True,
                text=True,
                timeout=max(5, ascend_probe_timeout_sec),
                env=env,
            )
            if result.returncode == 0:
                return True, None

            reason = (result.stderr or result.stdout or "Ascend 功能探针失败").strip()
            return False, reason.splitlines()[-1] if reason else "Ascend 功能探针失败"
        except subprocess.TimeoutExpired:
            return False, f"Ascend 功能探针超时（>{max(5, ascend_probe_timeout_sec)}s）"
        except Exception as exc:  # pragma: no cover - hardware/runtime dependent
            return False, str(exc)

    if not cpu:
        # Try Ascend first on Ascend-capable machines to avoid in-process torch
        # import side-effects before torch_npu initialization.
        if "✅" in detect_ascend():
            ascend_ok, ascend_reason = _ascend_functional_ready()
            if ascend_ok:
                backend_kind = "ascend"
                console.print("[green]✅ Ascend NPU detected[/green]")
            else:
                console.print(
                    "[yellow]⚠️  Ascend 检测到但功能验证失败，继续尝试 CUDA/CPU。[/yellow]"
                )
                if ascend_reason:
                    console.print(f"[dim]Ascend 失败原因: {ascend_reason}[/dim]")

        # Try CUDA when Ascend is unavailable/not ready.
        if backend_kind == "cpu":
            try:
                import torch

                if torch.cuda.is_available() and torch.cuda.device_count() > 0:
                    torch.zeros(1).cuda()  # Test CUDA actually works
                    backend_kind = "cuda"
                    gpu_name = torch.cuda.get_device_name(0)
                    console.print(f"[green]✅ CUDA detected:[/green] {gpu_name}")
            except Exception as exc:
                torch_probe_error = exc

    if torch_probe_error is not None:
        raise RuntimeError(
            "PyTorch failed to initialize during backend detection. "
            "Please verify torch/torch_npu installation compatibility. "
            f"Original error: {torch_probe_error}"
        ) from torch_probe_error

    if backend_kind == "cpu":
        console.print("[yellow]ℹ️  Using CPU backend[/yellow]")

    try:
        import asyncio

        hello_succeeded = False

        # Use a small but capable model for quick test
        model_name = "Qwen/Qwen2.5-0.5B-Instruct"
        prompt = "Hello, sageLLM!"

        console.print(f"[cyan]📦 Model:[/cyan] {model_name}")
        console.print(f"[cyan]💬 Prompt:[/cyan] {prompt}")
        console.print(
            f"[cyan]🔧 Backend:[/cyan] {backend_kind.upper() if backend_kind != 'cpu' else 'CPU'}\n"
        )

        async def _hello() -> None:
            nonlocal hello_succeeded
            from sagellm_core.llm_engine import LLMEngine, LLMEngineConfig
            from sagellm_protocol.sampling import SamplingParams

            provider_name = backend_kind

            config = LLMEngineConfig(
                model_path=model_name,
                backend_type=provider_name,
                max_new_tokens=30,
            )
            engine = LLMEngine(config)

            console.print("[dim]Loading model (first run downloads ~1GB)...[/dim]")
            await engine.start()

            try:
                sampling_params = SamplingParams(
                    max_tokens=20,
                    temperature=0.7,
                )

                response = await engine.generate(prompt, sampling_params=sampling_params)

                console.print("\n[green]✅ Inference successful![/green]")
                console.print(f"[cyan]📝 Output:[/cyan] {response.output_text[:100]}")
                console.print("\n[blue]📊 Metrics:[/blue]")
                console.print(f"   TTFT: {response.metrics.ttft_ms:.1f} ms")
                console.print(f"   Throughput: {response.metrics.throughput_tps:.1f} tokens/s")
                console.print("\n[bold green]🎉 sageLLM is working! Try:[/bold green]")
                console.print("   [cyan]sagellm run -p 'Your prompt here'[/cyan]")
                console.print("   [cyan]sagellm serve[/cyan]  # Start API server")
                console.print()
                hello_succeeded = True
            finally:
                await engine.stop()

        asyncio.run(_hello())
        if hello_succeeded and os.environ.get("SAGELLM_HELLO_FORCE_CLEAN_EXIT", "1") == "1":
            os._exit(0)

    except ImportError as e:
        console.print(f"[red]❌ Import error: {e}[/red]")
        console.print("\n[yellow]💡 Install the right PyTorch backend first:[/yellow]")
        console.print("   [cyan]sagellm install cuda[/cyan]    # NVIDIA GPU")
        console.print("   [cyan]sagellm install ascend[/cyan]  # Huawei Ascend NPU")
        console.print("   [cyan]sagellm install cpu[/cyan]     # CPU only (no GPU)")
        _exit_with_code(1)
    except Exception as e:
        console.print(f"[red]❌ Error: {e}[/red]")
        import traceback

        console.print(f"[dim]{traceback.format_exc()}[/dim]")
        _exit_with_code(1)
