"""Ascend diagnostics and repair command group."""

from __future__ import annotations

import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

import click

from ..utils.console import get_console


@dataclass
class AscendLayout:
    """Discovered Ascend toolkit layout."""

    root: Path
    toolkit_home: Path | None
    cann_link: Path
    latest_link: Path


def _discover_toolkit_home() -> Path | None:
    toolkit_root = Path("/usr/local/Ascend/ascend-toolkit")
    if not toolkit_root.is_dir():
        return None

    runtime_env = os.environ.get("ASCEND_RUNTIME_PATH")
    if runtime_env:
        runtime_path = Path(runtime_env)
        if runtime_path.is_dir() and runtime_path.name == "runtime":
            return runtime_path.parent

    for env_key in ("ASCEND_HOME", "ASCEND_HOME_PATH"):
        env_value = os.environ.get(env_key)
        if not env_value:
            continue
        env_path = Path(env_value)
        if env_path.is_dir() and (env_path / "opp").is_dir():
            return env_path

    latest_home = toolkit_root / "latest"
    if latest_home.is_dir() and (latest_home / "opp").is_dir():
        return latest_home

    candidates = sorted(
        [p for p in toolkit_root.iterdir() if p.is_dir() and p.name != "latest"], reverse=True
    )
    _arch_names = ("aarch64-linux", "x86_64-linux", "arm64-linux")
    for candidate in candidates:
        # Standard layout: opp/ + runtime/
        if (candidate / "opp").is_dir() and (candidate / "runtime").is_dir():
            return candidate
        # Arch-specific layout: <arch>/lib64/ (no top-level opp/)
        if any((candidate / arch / "lib64").is_dir() for arch in _arch_names):
            return candidate

    return None


def _collect_layout() -> AscendLayout:
    root = Path("/usr/local/Ascend")
    return AscendLayout(
        root=root,
        toolkit_home=_discover_toolkit_home(),
        cann_link=root / "cann",
        latest_link=root / "ascend-toolkit" / "latest",
    )


def _symlink_target(path: Path) -> str:
    if not path.exists() and not path.is_symlink():
        return "<missing>"
    if path.is_symlink():
        try:
            return str(path.resolve())
        except OSError:
            return "<broken-symlink>"
    return str(path)


def _repair_symlink(link_path: Path, target_path: Path) -> tuple[bool, str | None]:
    try:
        if link_path.is_symlink() or link_path.exists():
            link_path.unlink()
        link_path.symlink_to(target_path)
        return True, None
    except PermissionError as exc:
        return False, f"权限不足: {exc}"
    except OSError as exc:
        return False, str(exc)


def _apply_runtime_env(toolkit_home: Path) -> None:
    runtime_path = toolkit_home / "runtime"
    opp_path = toolkit_home / "opp"
    aicpu_fallback = toolkit_home / "aarch64-linux" / "lib64" / "plugin" / "opskernel"

    os.environ["ASCEND_HOME"] = str(toolkit_home)
    os.environ["ASCEND_HOME_PATH"] = str(toolkit_home)
    if runtime_path.is_dir():
        os.environ["ASCEND_RUNTIME_PATH"] = str(runtime_path)
    if opp_path.is_dir():
        os.environ["ASCEND_OPP_PATH"] = str(opp_path)
    # ASCEND_AICPU_PATH: prefer arch-specific opskernel dir, fall back to toolkit root
    if aicpu_fallback.is_dir():
        os.environ["ASCEND_AICPU_PATH"] = str(aicpu_fallback)
    else:
        # Check aicpu_kernels subdir (standard CANN >= 7.x layout) then fall back
        aicpu_kernels = toolkit_home / "aicpu_kernels"
        os.environ["ASCEND_AICPU_PATH"] = (
            str(aicpu_kernels) if aicpu_kernels.is_dir() else str(toolkit_home)
        )

    ld_prefixes = [
        toolkit_home / "fwkacllib" / "lib64",
        Path("/usr/local/Ascend/driver/lib64"),
        toolkit_home / "hccl" / "lib64",
        runtime_path / "lib64",
        toolkit_home / "lib64",
    ]
    existing = os.environ.get("LD_LIBRARY_PATH", "")
    merged: list[str] = []
    seen: set[str] = set()
    for prefix in ld_prefixes:
        prefix_str = str(prefix)
        if prefix.is_dir() and prefix_str not in seen:
            merged.append(prefix_str)
            seen.add(prefix_str)
    for segment in existing.split(":"):
        if segment and segment not in seen:
            merged.append(segment)
            seen.add(segment)
    os.environ["LD_LIBRARY_PATH"] = ":".join(merged)


def _run_functional_probe(timeout_sec: int = 60) -> tuple[bool, str]:
    probe_script = (
        "import os;"
        "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD']='0';"
        "import torch, torch_npu;"
        "print(f'npu_available={torch.npu.is_available()}');"
        "torch.npu.set_device('npu:0');"
        "x=torch.ones(1, device='npu');"
        "print((x+1).cpu().tolist())"
    )
    try:
        result = subprocess.run(
            [sys.executable, "-c", probe_script],
            capture_output=True,
            text=True,
            timeout=timeout_sec,
            env=os.environ.copy(),
        )
    except subprocess.TimeoutExpired:
        return False, "Ascend 功能探针超时"

    if result.returncode == 0:
        return True, (result.stdout or "probe ok").strip()

    details = (result.stderr or result.stdout or "probe failed").strip()
    tail = details.splitlines()[-1] if details else "probe failed"
    return False, tail


def _sudo_non_interactive_available() -> bool:
    try:
        result = subprocess.run(
            ["sudo", "-n", "true"],
            capture_output=True,
            text=True,
            timeout=5,
        )
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return False
    return result.returncode == 0


def _build_repair_script(toolkit_home: Path) -> str:
    return "\n".join(
        [
            "#!/usr/bin/env bash",
            "set -euo pipefail",
            "",
            f'TARGET_TOOLKIT="{toolkit_home}"',
            "",
            "echo '[1/4] Fix Ascend symlinks'",
            'sudo ln -sfn "${TARGET_TOOLKIT}" /usr/local/Ascend/cann',
            "sudo ln -sfn /usr/local/Ascend/cann /usr/local/Ascend/ascend-toolkit/latest",
            "",
            "echo '[2/4] Export runtime env for current shell'",
            'export ASCEND_HOME="${TARGET_TOOLKIT}"',
            'export ASCEND_HOME_PATH="${TARGET_TOOLKIT}"',
            'export ASCEND_RUNTIME_PATH="${TARGET_TOOLKIT}/runtime"',
            'export ASCEND_OPP_PATH="${TARGET_TOOLKIT}/opp"',
            'if [ -d "${TARGET_TOOLKIT}/aarch64-linux/lib64/plugin/opskernel" ]; then',
            '  export ASCEND_AICPU_PATH="${TARGET_TOOLKIT}/aarch64-linux/lib64/plugin/opskernel"',
            "fi",
            'export LD_LIBRARY_PATH="${TARGET_TOOLKIT}/fwkacllib/lib64:/usr/local/Ascend/driver/lib64:${TARGET_TOOLKIT}/hccl/lib64:${TARGET_TOOLKIT}/runtime/lib64:${TARGET_TOOLKIT}/lib64:${LD_LIBRARY_PATH:-}"',
            "",
            "echo '[3/4] Verify torch_npu functional probe'",
            "python - <<'PY'",
            "import torch, torch_npu",
            "print('npu_available=', torch.npu.is_available())",
            "torch.npu.set_device('npu:0')",
            "x=torch.ones(1, device='npu')",
            "print('tensor_ok=', (x+1).cpu().tolist())",
            "PY",
            "",
            "echo '[4/4] Run integrated CLI check'",
            "sagellm ascend check",
        ]
    )


@click.group()
def ascend() -> None:
    """Ascend environment diagnostics and repair."""


@ascend.command("check")
def ascend_check() -> None:
    """Check Ascend symlink/env/runtime health."""
    console = get_console()
    layout = _collect_layout()

    console.print("\n[bold blue]🔍 Ascend Health Check[/bold blue]")
    console.print(f"  root: {layout.root}")
    console.print(f"  cann: {_symlink_target(layout.cann_link)}")
    console.print(f"  latest: {_symlink_target(layout.latest_link)}")
    console.print(f"  toolkit_home: {layout.toolkit_home or '<missing>'}")

    if layout.toolkit_home is None:
        console.print("[red]❌ 未发现可用 Ascend toolkit 目录[/red]")
        raise SystemExit(1)

    _apply_runtime_env(layout.toolkit_home)
    ok, details = _run_functional_probe()
    if ok:
        console.print("[green]✅ Ascend 功能探针通过[/green]")
        console.print(f"[dim]{details}[/dim]")
    else:
        console.print("[yellow]⚠️ Ascend 功能探针失败[/yellow]")
        console.print(f"[dim]{details}[/dim]")
        raise SystemExit(2)


@ascend.command("fix")
@click.option("--run-probe/--no-probe", default=True, help="Run functional probe after repair")
@click.option(
    "--oneclick/--no-oneclick",
    default=False,
    help="Fallback to scripts/ascend-oneclick.sh --skip-doctor when probe still fails",
)
def ascend_fix(run_probe: bool, oneclick: bool) -> None:
    """Repair common Ascend path/env issues."""
    console = get_console()
    layout = _collect_layout()

    if layout.toolkit_home is None:
        console.print("[red]❌ 未找到 toolkit_home，无法自动修复[/red]")
        raise SystemExit(1)

    console.print("\n[bold blue]🛠️ Ascend Auto Fix[/bold blue]")
    console.print(f"  target toolkit: {layout.toolkit_home}")

    repaired = True
    msg: str | None = None

    changed, err = _repair_symlink(layout.cann_link, layout.toolkit_home)
    if not changed and err:
        repaired = False
        msg = err

    changed_latest, err_latest = _repair_symlink(layout.latest_link, layout.cann_link)
    if not changed_latest and err_latest:
        repaired = False
        msg = err_latest if msg is None else f"{msg}; {err_latest}"

    _apply_runtime_env(layout.toolkit_home)

    if repaired:
        console.print("[green]✅ 软链与运行时环境已更新[/green]")
    else:
        console.print("[yellow]⚠️ 部分软链修复失败[/yellow]")
        if msg:
            console.print(f"[dim]{msg}[/dim]")
        console.print(
            "[yellow]💡 需要更高权限时，请执行: sudo ln -sfn "
            f"{layout.toolkit_home} /usr/local/Ascend/cann[/yellow]"
        )

    if run_probe:
        ok, details = _run_functional_probe()
        if ok:
            console.print("[green]✅ Ascend 功能探针通过[/green]")
            console.print(f"[dim]{details}[/dim]")
            return

        console.print("[yellow]⚠️ 修复后功能探针仍失败[/yellow]")
        console.print(f"[dim]{details}[/dim]")

        if oneclick:
            script = Path(__file__).resolve().parents[4] / "scripts" / "ascend-oneclick.sh"
            if script.is_file():
                console.print("[cyan]↪ 继续执行 ascend-oneclick 深度修复...[/cyan]")
                proc = subprocess.run(["bash", str(script), "--skip-doctor"], env=os.environ.copy())
                if proc.returncode != 0:
                    raise SystemExit(proc.returncode)
                return

        raise SystemExit(2)


@ascend.command("doctor")
@click.option(
    "--print-script/--no-print-script",
    default=True,
    help="Print suggested sudo repair script",
)
@click.option(
    "--write-script",
    type=click.Path(path_type=Path),
    default=Path.home() / ".sagellm" / "ascend_repair.sh",
    show_default=True,
    help="Write suggested repair script to file",
)
def ascend_doctor(print_script: bool, write_script: Path) -> None:
    """Show diagnostic summary and generate a sudo repair script."""
    console = get_console()
    layout = _collect_layout()

    console.print("\n[bold blue]🩺 Ascend Doctor[/bold blue]")
    console.print(f"  root: {layout.root}")
    console.print(f"  cann: {_symlink_target(layout.cann_link)}")
    console.print(f"  latest: {_symlink_target(layout.latest_link)}")
    console.print(f"  toolkit_home: {layout.toolkit_home or '<missing>'}")
    console.print(f"  sudo -n available: {'yes' if _sudo_non_interactive_available() else 'no'}")

    if layout.toolkit_home is None:
        console.print("[red]❌ 未发现可用 toolkit_home，无法生成修复脚本[/red]")
        raise SystemExit(1)

    _apply_runtime_env(layout.toolkit_home)
    ok, details = _run_functional_probe(timeout_sec=30)
    if ok:
        console.print("[green]✅ 当前 Ascend 功能探针通过[/green]")
    else:
        console.print("[yellow]⚠️ 当前 Ascend 功能探针失败[/yellow]")
        console.print(f"[dim]{details}[/dim]")

    script_content = _build_repair_script(layout.toolkit_home)
    write_script.parent.mkdir(parents=True, exist_ok=True)
    write_script.write_text(script_content + "\n", encoding="utf-8")
    write_script.chmod(0o750)
    console.print(f"[green]✅ 修复脚本已生成: {write_script}[/green]")

    if print_script:
        console.print("\n[bold]Suggested repair script:[/bold]")
        console.print(script_content)
