"""orbit doctor — environment health check."""

from __future__ import annotations

import importlib
import os
import shutil
import sys

from rich.panel import Panel
from rich.table import Table

from orbit.utils.display import console


def _check_python() -> tuple[str, str, str]:
    """Check Python version."""
    v = sys.version_info
    ver = f"{v.major}.{v.minor}.{v.micro}"
    if v >= (3, 10):
        return "ok", f"Python {ver}", ""
    return "bad", f"Python {ver}", "Requires Python 3.10+"


def _check_package(name: str, import_name: str | None = None, min_version: str | None = None) -> tuple[str, str, str]:
    """Check if a package is installed."""
    mod_name = import_name or name
    try:
        mod = importlib.import_module(mod_name)
        version = getattr(mod, "__version__", getattr(mod, "VERSION", "installed"))
        return "ok", f"{name} {version}", ""
    except ImportError:
        return "bad", f"{name} not installed", f"pip install {name}"


def _check_gpu() -> tuple[str, str, str]:
    """Check GPU availability."""
    try:
        import torch

        if torch.cuda.is_available():
            name = torch.cuda.get_device_name(0)
            mem = torch.cuda.get_device_properties(0).total_mem / (1024**3)
            return "ok", f"GPU available: {name} ({mem:.0f}GB)", ""
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "ok", "Apple Silicon MPS available", ""
        return "warn", "No GPU detected (CPU only)", "GPU recommended for --proxy and embeddings"
    except ImportError:
        return "warn", "torch not installed (GPU check skipped)", "pip install torch"


def _check_hf_hub() -> tuple[str, str, str]:
    """Check HuggingFace Hub accessibility."""
    try:
        from huggingface_hub import HfApi

        api = HfApi()
        # Quick connectivity check
        api.list_models(limit=1)
        return "ok", "HuggingFace Hub accessible", ""
    except Exception as e:
        return "warn", f"HuggingFace Hub: {e}", "Check internet connection"


def _check_env_var(var: str, purpose: str) -> tuple[str, str, str]:
    """Check if an environment variable is set."""
    if os.environ.get(var):
        return "ok", f"{var} is set", ""
    return "warn", f"{var} not set", f"Needed for {purpose}"


def _check_lerobot() -> tuple[str, str, str]:
    """Check lerobot installation."""
    try:
        import lerobot

        ver = getattr(lerobot, "__version__", "installed")
        return "ok", f"lerobot {ver}", ""
    except ImportError:
        return "warn", "lerobot not installed", "pip install lerobot"


def _check_tool(name: str, purpose: str) -> tuple[str, str, str]:
    """Check if a CLI tool is on PATH."""
    if shutil.which(name):
        return "ok", f"{name} found on PATH", ""
    return "warn", f"{name} not found", f"Needed for {purpose}"


def run_doctor() -> dict:
    """Run all health checks and return results."""
    checks: list[tuple[str, str, str, str]] = []  # (category, status, message, fix)

    # Core environment
    status, msg, fix = _check_python()
    checks.append(("Core", status, msg, fix))

    # Core packages
    for pkg, imp in [("numpy", None), ("pandas", None), ("scipy", None),
                     ("rich", None), ("click", None), ("huggingface_hub", "huggingface_hub"),
                     ("pyarrow", None), ("scikit-learn", "sklearn")]:
        status, msg, fix = _check_package(pkg, imp)
        checks.append(("Core", status, msg, fix))

    # HuggingFace connectivity
    status, msg, fix = _check_hf_hub()
    checks.append(("Connectivity", status, msg, fix))

    # Optional: vision stack
    for pkg, imp in [("torch", None), ("transformers", None),
                     ("opencv-python", "cv2"), ("decord", None), ("Pillow", "PIL")]:
        status, msg, fix = _check_package(pkg, imp)
        if status == "bad":
            status = "warn"
        checks.append(("Vision (optional)", status, msg, fix))

    # Optional: VLM / AI
    status, msg, fix = _check_package("google-generativeai", "google.generativeai")
    if status == "bad":
        status = "warn"
    checks.append(("AI (optional)", status, msg, fix))
    status, msg, fix = _check_env_var("GOOGLE_API_KEY", "--ai diagnosis and VLM assessment")
    checks.append(("AI (optional)", status, msg, fix))

    # GPU
    status, msg, fix = _check_gpu()
    checks.append(("Hardware", status, msg, fix))

    # Ecosystem tools
    status, msg, fix = _check_lerobot()
    checks.append(("Ecosystem", status, msg, fix))

    for pkg, imp in [("mujoco", None), ("h5py", None)]:
        status, msg, fix = _check_package(pkg, imp)
        if status == "bad":
            status = "warn"
        checks.append(("Ecosystem (optional)", status, msg, fix))

    return checks


def print_doctor_report(checks: list) -> None:
    """Print the doctor report with Rich formatting."""
    icons = {"ok": "[green]\u2713[/green]", "warn": "[yellow]\u26a0[/yellow]", "bad": "[red]\u2717[/red]"}

    table = Table(show_header=False, box=None, padding=(0, 1))
    table.add_column("status", width=3)
    table.add_column("message")
    table.add_column("fix", style="dim")

    prev_cat = None
    ok_count = warn_count = bad_count = 0
    for cat, status, msg, fix in checks:
        if cat != prev_cat:
            if prev_cat is not None:
                table.add_row("", "", "")
            table.add_row("", f"[bold]{cat}[/bold]", "")
            prev_cat = cat
        table.add_row(icons[status], msg, fix)
        if status == "ok":
            ok_count += 1
        elif status == "warn":
            warn_count += 1
        else:
            bad_count += 1

    console.print()
    console.print(Panel("[bold]ORBIT Doctor[/bold] \u2014 Environment Health Check", border_style="cyan"))
    console.print(table)
    console.print()

    summary_parts = [f"[green]{ok_count} passed[/green]"]
    if warn_count:
        summary_parts.append(f"[yellow]{warn_count} warnings[/yellow]")
    if bad_count:
        summary_parts.append(f"[red]{bad_count} errors[/red]")
    sep = " \u00b7 "
    console.print(f"  {sep.join(summary_parts)}")

    if bad_count:
        console.print("\n  [bold red]Some required checks failed.[/bold red] Fix errors above to use ORBIT.")
    elif warn_count:
        console.print("\n  [dim]Optional features may be unavailable. Install missing packages for full functionality.[/dim]")
        console.print("  [dim]Quick fix: pip install orbit-robotics[all][/dim]")
    else:
        console.print("\n  [bold green]All checks passed![/bold green] ORBIT is fully configured.")
    console.print()
