"""Interactive setup wizard for model recommendation."""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path

import click

from ..hardware.ascend import check_ascend_available
from ..hardware.cuda import check_cuda_available
from ..hardware.domestic import (
    check_haiguang_available,
    check_kunlun_available,
    check_muxi_available,
)
from ..utils.console import get_console


@dataclass(frozen=True)
class MachineProfile:
    """Detected machine capabilities."""

    cpu_cores: int
    ram_gb: float
    cuda_ready: bool
    cuda_name: str
    cuda_vram_gb: float
    ascend_ready: bool
    kunlun_ready: bool
    haiguang_ready: bool
    muxi_ready: bool


@dataclass(frozen=True)
class ModelCandidate:
    """Candidate model entry for recommendations."""

    model_id: str
    suggested_backends: tuple[str, ...]
    min_ram_gb: float
    min_vram_gb: float
    note: str


CANDIDATES: tuple[ModelCandidate, ...] = (
    ModelCandidate(
        model_id="Qwen/Qwen2.5-0.5B-Instruct",
        suggested_backends=("cpu", "cuda", "ascend", "kunlun", "haiguang", "muxi"),
        min_ram_gb=6,
        min_vram_gb=2,
        note="轻量模型，启动快，适合首次验证",
    ),
    ModelCandidate(
        model_id="Qwen/Qwen2.5-1.5B-Instruct",
        suggested_backends=("cuda", "ascend", "haiguang", "muxi", "cpu"),
        min_ram_gb=12,
        min_vram_gb=6,
        note="通用推荐，质量/速度平衡",
    ),
    ModelCandidate(
        model_id="Qwen/Qwen2.5-3B-Instruct",
        suggested_backends=("cuda", "ascend", "haiguang", "muxi"),
        min_ram_gb=20,
        min_vram_gb=10,
        note="效果更好，建议有独立加速卡",
    ),
    ModelCandidate(
        model_id="Qwen/Qwen2.5-7B-Instruct",
        suggested_backends=("cuda", "ascend", "haiguang"),
        min_ram_gb=32,
        min_vram_gb=18,
        note="高质量候选，建议中高端显存",
    ),
)


BACKEND_PRIORITY: tuple[str, ...] = ("ascend", "kunlun", "haiguang", "muxi", "cuda", "cpu")


def _read_ram_gb() -> float:
    meminfo = Path("/proc/meminfo")
    if not meminfo.exists():
        return 0.0
    for line in meminfo.read_text(encoding="utf-8").splitlines():
        if line.startswith("MemTotal:"):
            parts = line.split()
            if len(parts) >= 2 and parts[1].isdigit():
                mem_kb = int(parts[1])
                return round(mem_kb / (1024 * 1024), 1)
    return 0.0


def _detect_profile() -> MachineProfile:
    cuda_ready, _ = check_cuda_available()
    ascend_ready, _ = check_ascend_available()
    kunlun_ready, _ = check_kunlun_available()
    haiguang_ready, _ = check_haiguang_available()
    muxi_ready, _ = check_muxi_available()

    cuda_name = "-"
    cuda_vram_gb = 0.0

    if cuda_ready:
        try:
            import torch

            cuda_name = str(torch.cuda.get_device_name(0))
            total_mem = torch.cuda.get_device_properties(0).total_memory
            cuda_vram_gb = round(total_mem / (1024**3), 1)
        except Exception:
            cuda_name = "CUDA device"
            cuda_vram_gb = 0.0

    return MachineProfile(
        cpu_cores=max(1, int(os.cpu_count() or 1)),
        ram_gb=_read_ram_gb(),
        cuda_ready=cuda_ready,
        cuda_name=cuda_name,
        cuda_vram_gb=cuda_vram_gb,
        ascend_ready=ascend_ready,
        kunlun_ready=kunlun_ready,
        haiguang_ready=haiguang_ready,
        muxi_ready=muxi_ready,
    )


def _available_backends(profile: MachineProfile) -> list[str]:
    available: list[str] = ["cpu"]
    if profile.cuda_ready:
        available.append("cuda")
    if profile.ascend_ready:
        available.append("ascend")
    if profile.kunlun_ready:
        available.append("kunlun")
    if profile.haiguang_ready:
        available.append("haiguang")
    if profile.muxi_ready:
        available.append("muxi")
    return available


def _preferred_backend(profile: MachineProfile) -> str:
    available = set(_available_backends(profile))
    for backend in BACKEND_PRIORITY:
        if backend in available:
            return backend
    return "cpu"


def _recommend(profile: MachineProfile) -> list[tuple[ModelCandidate, str, str]]:
    preferred = _preferred_backend(profile)
    available = set(_available_backends(profile))
    picks: list[tuple[ModelCandidate, str, str]] = []

    for item in CANDIDATES:
        backend = preferred if preferred in item.suggested_backends else "cpu"
        if backend not in available:
            continue

        if profile.ram_gb < item.min_ram_gb:
            continue
        if backend != "cpu" and profile.cuda_ready and item.min_vram_gb > 0:
            if profile.cuda_vram_gb and profile.cuda_vram_gb < item.min_vram_gb:
                continue

        reason = f"后端 {backend}，内存≥{item.min_ram_gb}GB"
        if backend != "cpu" and item.min_vram_gb > 0:
            reason = f"后端 {backend}，显存≥{item.min_vram_gb}GB"
        picks.append((item, backend, reason))

    if not picks:
        fallback = CANDIDATES[0]
        picks.append((fallback, "cpu", "资源较紧张，建议先从轻量模型开始"))

    return picks


def _render_profile(profile: MachineProfile) -> None:
    from rich.table import Table

    console = get_console()
    table = Table(title="🧪 机器检测结果", show_header=True, header_style="bold magenta")
    table.add_column("项目", style="cyan", width=18)
    table.add_column("结果", style="white")

    table.add_row("CPU", f"{profile.cpu_cores} cores")
    table.add_row("内存", f"{profile.ram_gb} GB")
    table.add_row("CUDA", "✅" if profile.cuda_ready else "❌")
    if profile.cuda_ready:
        table.add_row("GPU", f"{profile.cuda_name} ({profile.cuda_vram_gb} GB)")
    table.add_row("Ascend", "✅" if profile.ascend_ready else "❌")
    table.add_row("Kunlun", "✅" if profile.kunlun_ready else "❌")
    table.add_row("Haiguang", "✅" if profile.haiguang_ready else "❌")
    table.add_row("Muxi", "✅" if profile.muxi_ready else "❌")

    console.print(table)


def _render_recommendations(recommended: list[tuple[ModelCandidate, str, str]]) -> None:
    from rich.table import Table

    console = get_console()
    table = Table(title="🤖 推荐模型候选", show_header=True, header_style="bold green")
    table.add_column("编号", width=6)
    table.add_column("模型")
    table.add_column("建议后端", width=10)
    table.add_column("说明")

    for idx, (item, backend, reason) in enumerate(recommended, start=1):
        table.add_row(str(idx), item.model_id, backend, f"{item.note}；{reason}")

    console.print(table)


def _print_commands(choices: list[tuple[ModelCandidate, str, str]]) -> None:
    console = get_console()
    console.print("\n[bold blue]可直接执行的命令：[/bold blue]")
    for item, backend, _ in choices:
        console.print(
            f"  [cyan]sagellm run -m {item.model_id} --backend {backend} -p '你好，请做个自我介绍'[/cyan]"
        )
    if len(choices) == 1:
        item, backend, _ = choices[0]
        console.print("\n[dim]如需启动服务接口：[/dim]")
        console.print(
            f"  [dim]sagellm serve --model {item.model_id} --backend {backend} --port 8000[/dim]"
        )


def _select_choices(
    recommended: list[tuple[ModelCandidate, str, str]],
) -> list[tuple[ModelCandidate, str, str]]:
    raw = click.prompt("请输入要试跑的模型编号（支持逗号分隔）", default="1").strip()
    picked_indices: list[int] = []
    for part in raw.split(","):
        token = part.strip()
        if not token:
            continue
        if token.isdigit():
            idx = int(token)
            if 1 <= idx <= len(recommended):
                picked_indices.append(idx)

    if not picked_indices:
        return [recommended[0]]

    unique = sorted(set(picked_indices))
    return [recommended[i - 1] for i in unique]


@click.command(name="setup")
def setup() -> None:
    """交互式向导：检测机器并推荐可运行的大模型。"""
    console = get_console()
    console.print("\n[bold blue]🧭 sagellm Setup Wizard[/bold blue]")
    console.print("[dim]根据你的机器配置推荐可运行模型，并生成一键命令。[/dim]\n")

    profile = _detect_profile()

    while True:
        _render_profile(profile)
        recommended = _recommend(profile)
        _render_recommendations(recommended)

        console.print("\n[bold]菜单[/bold]")
        console.print("  [cyan]1[/cyan] 自动使用 Top-1 推荐")
        console.print("  [cyan]2[/cyan] 手动选择多个候选")
        console.print("  [cyan]3[/cyan] 重新检测机器")
        console.print("  [cyan]0[/cyan] 退出")

        action = click.prompt("请选择", default="1").strip()
        if action == "0":
            console.print("[yellow]已退出 setup 向导。[/yellow]\n")
            return
        if action == "3":
            profile = _detect_profile()
            console.print("[green]已重新检测硬件。[/green]\n")
            continue
        if action == "2":
            choices = _select_choices(recommended)
        else:
            choices = [recommended[0]]

        _print_commands(choices)
        console.print()

        if click.confirm("是否继续在向导里选择其他模型", default=False):
            continue
        return
