"""Interactive first-run setup wizard — mining-client style."""

from __future__ import annotations

import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import TYPE_CHECKING

from .config import DEFAULTS, _DEFAULT_DOWNLOAD_DIR, _DEFAULT_TRACKER_URL, _DEFAULT_WALLET_PATH
from .model_catalog import CURATED_MODELS, ModelPreset, browse_hf_hub

if TYPE_CHECKING:
    pass

_HEADER = """\
╔══════════════════════════════════════════════════════════════════╗
║              meshnet-node  v0.1.0                                ║
║        Distributed AI Inference — Node Setup                     ║
╚══════════════════════════════════════════════════════════════════╝
"""

_QUANT_LABELS = {"nf4": "NF4 (4-bit)", "int8": "INT8 (8-bit)", "bf16": "BF16 (full)"}


def _ask(prompt: str, default: str = "", validator=None) -> str:
    """Prompt user and return answer. Returns default on empty input or EOF."""
    display = f"{prompt} [{default}]: " if default else f"{prompt}: "
    while True:
        try:
            raw = input(display).strip()
        except (EOFError, KeyboardInterrupt):
            print()
            raise KeyboardInterrupt
        value = raw or default
        if validator is None or validator(value):
            return value
        # validator returned error string
        print(f"  ✗ {validator(value)}")


def _ask_int(prompt: str, default: int, lo: int, hi: int) -> int:
    def validate(s: str) -> bool | str:
        try:
            v = int(s)
        except ValueError:
            return "Please enter a number."
        if not (lo <= v <= hi):
            return f"Please enter a number between {lo} and {hi}."
        return True

    while True:
        raw = _ask(prompt, str(default))
        try:
            v = int(raw)
            if lo <= v <= hi:
                return v
        except ValueError:
            pass
        print(f"  ✗ Enter a number between {lo} and {hi}.")


def _ask_yn(prompt: str, default: bool = True) -> bool:
    hint = "Y/n" if default else "y/N"
    raw = _ask(f"{prompt} [{hint}]").lower()
    if not raw:
        return default
    return raw.startswith("y")


def _detect_gpus() -> list[dict]:
    """Return list of detected GPU dicts with name and vram_gb."""
    gpus: list[dict] = []
    try:
        import torch  # type: ignore[import]
        if torch.cuda.is_available():
            for i in range(torch.cuda.device_count()):
                props = torch.cuda.get_device_properties(i)
                gpus.append(
                    {
                        "index": i,
                        "name": props.name,
                        "vram_gb": props.total_memory / 1e9,
                        "backend": "cuda",
                    }
                )
    except ImportError:
        pass
    return gpus


def _total_vram_gb(gpus: list[dict]) -> float:
    return sum(g["vram_gb"] for g in gpus)


def _print_gpus(gpus: list[dict]) -> None:
    if not gpus:
        print("  ⚠ No CUDA GPU detected — running in CPU mode")
        print("    CPU inference is very slow. Consider a machine with an NVIDIA GPU.")
        return
    for g in gpus:
        vram = g["vram_gb"]
        print(f"  GPU {g['index']}: {g['name']}  {vram:.0f} GB VRAM  ✓")


def _print_model_table(gpus: list[dict], quant: str = "nf4") -> None:
    available_gb = _total_vram_gb(gpus)
    print()
    print(f" #   {'Model':<30} {'Layers':>6}   {'NF4':>6}   {'INT8':>6}   {'BF16':>6}")
    print(f" {'─'*4} {'─'*30} {'─'*6}   {'─'*6}   {'─'*6}   {'─'*6}")
    for i, m in enumerate(CURATED_MODELS, 1):
        fits_nf4 = "✓" if m.vram_nf4 <= available_gb else "✗"
        fits_int8 = "✓" if m.vram_int8 <= available_gb else "✗"
        fits_bf16 = "✓" if m.vram_bf16 <= available_gb else "✗"
        nf4_str = f"{fits_nf4}{m.vram_nf4:.0f}GB"
        int8_str = f"{fits_int8}{m.vram_int8:.0f}GB"
        bf16_str = f"{fits_bf16}{m.vram_bf16:.0f}GB"
        print(f" {i:<3} {m.name:<30} {m.num_layers:>6}   {nf4_str:>6}   {int8_str:>6}   {bf16_str:>6}")
        print(f"     {m.description}")
    idx = len(CURATED_MODELS) + 1
    print(f" {idx:<3} {'[Browse HuggingFace Hub...]':<30}")
    print()


def _browse_hf_interactive() -> str | None:
    """Show HF Hub top-20 and let user enter a repo ID. Returns repo ID or None to go back."""
    print("\nFetching top models from HuggingFace Hub...")
    try:
        models = browse_hf_hub(top_n=20)
    except RuntimeError as exc:
        print(f"  ✗ {exc}")
        return None

    print(f"\n {'#':<4} {'HuggingFace Repo':<50} Downloads")
    print(f" {'─'*4} {'─'*50} {'─'*10}")
    for i, m in enumerate(models, 1):
        dl = m["downloads"]
        dl_str = f"{dl/1e6:.1f}M" if dl >= 1_000_000 else f"{dl/1e3:.0f}k" if dl >= 1000 else str(dl)
        print(f" {i:<4} {m['repo']:<50} {dl_str}")

    print()
    raw = _ask(
        "Enter a number to select, or paste any HuggingFace repo ID (or press Enter to go back)",
        default="",
    )
    if not raw:
        return None
    try:
        idx = int(raw) - 1
        if 0 <= idx < len(models):
            return models[idx]["repo"]
    except ValueError:
        pass
    # Treat raw input as a repo ID
    if "/" in raw:
        return raw
    print("  ✗ Invalid input — please enter a number or a full repo ID like 'org/model-name'")
    return None


def _ask_quant(gpus: list[dict], model: ModelPreset | None) -> str:
    available_gb = _total_vram_gb(gpus)
    print("\nQuantization level:")
    options: list[tuple[str, str]] = []
    for quant, label in [("nf4", "NF4 4-bit"), ("int8", "INT8 8-bit"), ("bf16", "BF16 full precision")]:
        if model is not None:
            vram = model.vram_for_quant(quant)
            fits = "✓" if vram <= available_gb else "✗ insufficient VRAM"
            suffix = f"  ({vram:.0f} GB needed — {fits})"
        else:
            suffix = ""
        options.append((quant, f"{label}{suffix}"))

    for i, (_, label) in enumerate(options, 1):
        print(f"  {i}) {label}")

    # Recommend the best fitting quant
    if model is not None:
        rec = model.recommended_quant(available_gb)
        rec_idx = next((i for i, (q, _) in enumerate(options, 1) if q == rec), 1) if rec else 1
        default_idx = rec_idx
        print(f"  (Recommended: {rec.upper() if rec else 'NF4'} for your GPU)")
    else:
        default_idx = 1

    choice = _ask_int("Enter number", default_idx, 1, 3)
    return options[choice - 1][0]


def _validate_dir(path_str: str) -> bool | str:
    p = Path(path_str).expanduser()
    try:
        p.mkdir(parents=True, exist_ok=True)
        return True
    except OSError as exc:
        return f"Cannot create directory: {exc}"


def _validate_tracker(url: str) -> bool | str:
    if not url.startswith(("http://", "https://")):
        return "URL must start with http:// or https://"
    return True


def _ping_tracker(url: str) -> bool:
    """Return True if tracker responds to /health."""
    try:
        with urllib.request.urlopen(f"{url.rstrip('/')}/health", timeout=3):
            return True
    except Exception:
        return False


def run_wizard(config_path_override=None) -> dict:
    """Run the interactive setup wizard and return a config dict.

    Raises KeyboardInterrupt if user presses Ctrl-C.
    """
    print(_HEADER)

    # Step 1: GPU detection
    print("Detecting hardware...")
    gpus = _detect_gpus()
    _print_gpus(gpus)
    available_gb = _total_vram_gb(gpus)
    if available_gb == 0:
        available_gb = 9999  # CPU — don't filter models by VRAM

    # Step 2 & 3: Model selection
    print("\nSelect a model to serve:\n")
    selected_repo: str | None = None
    selected_preset: ModelPreset | None = None

    while selected_repo is None:
        _print_model_table(gpus)
        lo, hi = 1, len(CURATED_MODELS) + 1
        choice = _ask_int("Enter number", 1, lo, hi)
        if choice == len(CURATED_MODELS) + 1:
            repo = _browse_hf_interactive()
            if repo:
                selected_repo = repo
                selected_preset = None
        else:
            selected_preset = CURATED_MODELS[choice - 1]
            selected_repo = selected_preset.hf_repo
            if selected_preset.recommended_quant(available_gb) is None:
                print(
                    f"\n  ⚠ Warning: {selected_preset.name} requires at least "
                    f"{selected_preset.vram_nf4:.0f} GB VRAM at NF4 — even the smallest "
                    f"quantization may be too large for your GPU."
                )
                if not _ask_yn("Continue anyway?", default=False):
                    selected_repo = None
                    selected_preset = None

    print(f"\n  ✓ Selected: {selected_repo}")

    # Step 3b: Quantization
    quant = _ask_quant(gpus, selected_preset)
    print(f"  ✓ Quantization: {quant.upper()}")

    # Step 4: Download directory
    print()
    dl_dir = _ask(
        "Download directory",
        default=str(_DEFAULT_DOWNLOAD_DIR),
        validator=lambda v: _validate_dir(v) if v else "Directory is required.",
    )
    print(f"  ✓ Download dir: {dl_dir}")

    # Step 5: Tracker URL
    print()
    tracker_url = _DEFAULT_TRACKER_URL
    raw_tracker = _ask("Tracker URL", default=_DEFAULT_TRACKER_URL, validator=_validate_tracker)
    tracker_url = raw_tracker
    if _ping_tracker(tracker_url):
        print(f"  ✓ Tracker reachable: {tracker_url}")
    else:
        print(f"  ⚠ Tracker not reachable at {tracker_url} (will retry on start)")

    # Step 6: Wallet path
    print()
    wallet_path = _ask("Wallet path", default=_DEFAULT_WALLET_PATH)
    print(f"  ✓ Wallet: {wallet_path}")

    cfg = {
        "model_hf_repo": selected_repo,
        "model_name": selected_preset.name if selected_preset else selected_repo.split("/")[-1],
        "quantization": quant,
        "download_dir": dl_dir,
        "tracker_url": tracker_url,
        "wallet_path": wallet_path,
        "shard_start": None,
        "shard_end": None,
        "port": DEFAULTS["port"],
        "host": DEFAULTS["host"],
    }
    return cfg


def print_models_table(available_gb: float | None = None) -> None:
    """Print curated model table for `meshnet-node models`."""
    gpus: list[dict] = []
    if available_gb is None:
        gpus = _detect_gpus()
        available_gb = _total_vram_gb(gpus) or 9999
    else:
        gpus = [{"index": 0, "name": "GPU", "vram_gb": available_gb, "backend": "cuda"}]

    print(f"\n{'#':<4} {'Model':<32} {'HuggingFace Repo':<45} {'Layers':>6}  {'NF4':>8} {'INT8':>8} {'BF16':>8}")
    print(f"{'─'*4} {'─'*32} {'─'*45} {'─'*6}  {'─'*8} {'─'*8} {'─'*8}")
    for i, m in enumerate(CURATED_MODELS, 1):
        def _cell(vram: float) -> str:
            fits = "✓" if vram <= available_gb else "✗"
            return f"{fits}{vram:.0f}GB"

        print(
            f"{i:<4} {m.name:<32} {m.hf_repo:<45} {m.num_layers:>6}  "
            f"{_cell(m.vram_nf4):>8} {_cell(m.vram_int8):>8} {_cell(m.vram_bf16):>8}"
        )
    print()