"""
llmpm install <repo_id>

Downloads a model from HuggingFace Hub and registers it locally.

Examples:
  llmpm install bartowski/Llama-3.2-1B-Instruct-GGUF
  llmpm install bartowski/Llama-3.2-1B-Instruct-GGUF --quant Q5_K_M
  llmpm install meta-llama/Llama-3.2-1B-Instruct
  llmpm install meta-llama/Llama-3.2-1B-Instruct --file model.safetensors
"""

from __future__ import annotations

import sys
import time
from pathlib import Path

_IS_WIN = sys.platform == "win32"

import click

from llmpm import display
from llmpm.core import context, downloader, llmpm_json, registry
from llmpm.core.downloader import normalize_repo_id


@click.command("install")
@click.argument("repo_ids", nargs=-1)
@click.option(
    "--quant", "-q",
    default=None,
    help="For GGUF models: quantisation to download (e.g. Q4_K_M). "
         "Prompted interactively if omitted.",
)
@click.option(
    "--file", "-f", "specific_file",
    default=None,
    help="Download a specific file from the repo (overrides --quant).",
)
@click.option(
    "--no-interactive", is_flag=True,
    help="Never prompt; pick the best quantisation automatically.",
)
@click.option(
    "--save", "save", is_flag=True,
    help="Add the model(s) to llmpm.json dependencies after installing.",
)
def install(
    repo_ids: tuple[str, ...],
    quant: str | None,
    specific_file: str | None,
    no_interactive: bool,
    save: bool,
) -> None:
    """Download and install one or more models from HuggingFace Hub.

    When called with no arguments inside a project that has a llmpm.json,
    installs every model listed under ``dependencies`` (like ``npm install``).

    Use ``--save`` to install model(s) and record them in llmpm.json.
    """
    # ── no-arg: batch install from llmpm.json ────────────────────────────────
    if not repo_ids:
        deps = llmpm_json.dependencies()
        if not deps:
            if context.is_local():
                display.header("install")
                display.blank()
                display.info(
                    "No dependencies in [bold]llmpm.json[/].\n\n"
                    "  Add one:  [bold cyan]llmpm install <repo> --save[/]"
                )
            else:
                display.error(
                    "No model specified and no [bold]llmpm.json[/] found.\n\n"
                    "  Usage:  [bold cyan]llmpm install <repo>[/]\n"
                    "  Init a project first:  [bold cyan]llmpm init[/]"
                )
            display.blank()
            raise SystemExit(0)

        display.header(f"install  [dim](from llmpm.json · {len(deps)} model(s))[/]")
        display.blank()
        _show_local_mode_note()

        for dep_id in deps:
            _install_one(dep_id, quant=None, specific_file=None, no_interactive=True)
        raise SystemExit(0)

    # ── one or more models ────────────────────────────────────────────────────
    _show_local_mode_note()

    no_llmpm_json_warned = False
    for repo_id in repo_ids:
        resolved_id = _install_one(repo_id, quant, specific_file, no_interactive)

        if save:
            if llmpm_json.add_dependency(resolved_id):
                display.ok(f"Added [bold]{resolved_id}[/] to [bold]llmpm.json[/] dependencies")
                display.blank()
            elif not no_llmpm_json_warned:
                display.warn(
                    "[bold]--save[/] has no effect: no [bold]llmpm.json[/] found.\n"
                    "  Run [bold cyan]llmpm init[/] first to create one."
                )
                display.blank()
                no_llmpm_json_warned = True


# ── single-model install logic ────────────────────────────────────────────────

def _install_one(
    repo_id: str,
    quant: str | None,
    specific_file: str | None,
    no_interactive: bool,
) -> str:
    """Resolve, download, and register a single model. Returns the resolved repo_id."""
    repo_id = normalize_repo_id(repo_id)
    display.header(f"install {repo_id}")
    display.blank()

    display.step("Resolving model…")
    try:
        info = downloader.fetch_model_info(repo_id)
    except ValueError as exc:
        if "not found" not in str(exc).lower():
            display.error(str(exc))
            raise SystemExit(1) from None

        # ── Model not found: search for close matches ─────────────────────────
        display.warn(f"Model [bold]{repo_id}[/] not found.")
        display.blank()
        display.step(f"Searching for \"{repo_id}\"…")

        results = downloader.search_models(repo_id)
        if not results:
            display.error(
                f"No models found matching [bold]{repo_id}[/].\n\n"
                "  Browse models at: [bold cyan]https://huggingface.co/models[/]"
            )
            raise SystemExit(1) from None

        repo_id = _pick_from_search_results(results, repo_id)
        if repo_id is None:
            raise SystemExit(0)

        try:
            info = downloader.fetch_model_info(repo_id)
        except ValueError as exc2:
            display.error(str(exc2))
            raise SystemExit(1) from None

        display.header(f"install {repo_id}")
        display.blank()

    display.ok(f"[bold]{repo_id}[/]")
    display.blank()

    display.model_card(
        repo_id=repo_id,
        model_type=info.model_type,
        tags=info.tags,
        total_size=info.total_size,
        file_count=len(info.all_files),
    )
    display.blank()

    dest_dir = registry.model_install_path(repo_id)

    if specific_file:
        _install_single_file(info, repo_id, specific_file, dest_dir)
        return repo_id

    if info.is_gguf:
        _install_gguf(info, repo_id, quant, dest_dir, no_interactive)
    else:
        _install_transformer(info, repo_id, dest_dir)

    return repo_id


def _download_error_message(repo_id: str, exc: Exception) -> str:
    """Return a user-friendly error message for a download failure.

    For 401 / gated-repo errors, appends an HF_TOKEN hint.
    """
    msg = str(exc)
    is_auth = (
        "401" in msg
        or "gated" in msg.lower()
        or "restricted" in msg.lower()
        or "authentication" in msg.lower()
        or "log in" in msg.lower()
    )
    if is_auth:
        return (
            f"Download failed: access to [bold]{repo_id}[/] is restricted.\n\n"
            "  This is a gated model — you need to:\n"
            "    1. Accept the licence at "
            "[bold cyan]https://huggingface.co/[/bold cyan][bold]"
            f"{repo_id}[/bold]\n"
            "    2. Re-run with your HF token:\n\n"
            f"         [bold cyan]HF_TOKEN=<your_token> llmpm install {repo_id}[/]"
        )
    return f"Download failed: {exc}"


def _show_local_mode_note() -> None:
    """Print a one-line local-mode indicator when inside a project."""
    if context.is_local():
        home = context.get_home()
        display.info(f"local mode  [dim]{home}[/]")
        display.blank()


# ── HuggingFace search picker ─────────────────────────────────────────────────

def _fmt_downloads(n: int) -> str:
    """Format a download count as a short human-readable string."""
    if n >= 1_000_000:
        return f"{n / 1_000_000:.1f}M↓"
    if n >= 1_000:
        return f"{n / 1_000:.0f}K↓"
    return f"{n}↓"


def _pick_from_search_results(
    results: list[dict],
    original_query: str,
) -> str | None:
    """
    Present search results as a selection list and return the chosen repo_id.
    Returns None if the user cancels.
    """
    display.ok(f"Found {len(results)} model(s) matching \"{original_query}\":")
    display.blank()

    # Build labelled choices: "owner/model-name  (pipeline · 1.2M↓)"
    try:
        if _IS_WIN:
            raise RuntimeError("skip questionary on Windows")
        import questionary  # type: ignore  # pylint: disable=import-outside-toplevel
        from questionary import Choice  # pylint: disable=import-outside-toplevel

        choices = [
            Choice(
                title=(
                    f"{r['repo_id']:<55}"
                    f"  [dim]{r['pipeline_tag']:<22}"
                    f"  {_fmt_downloads(r['downloads'])}[/dim]"
                ),
                value=r["repo_id"],
            )
            for r in results
        ]

        selected = questionary.select(
            "Select a model to install:",
            choices=choices,
            style=questionary.Style([
                ("selected", "bold fg:green"),
                ("pointer", "bold fg:green"),
            ]),
        ).ask()

    except Exception:  # pylint: disable=broad-except
        # Fallback: numbered list + prompt
        for idx, r in enumerate(results, 1):
            tag = r["pipeline_tag"]
            dl = _fmt_downloads(r["downloads"])
            click.echo(f"  {idx:>2}.  {r['repo_id']:<55}  {tag:<22}  {dl}")
        display.blank()
        choice_strs = [str(i) for i in range(1, len(results) + 1)]
        raw = click.prompt(
            "  Select (number)",
            type=click.Choice(choice_strs),
        )
        selected = results[int(raw) - 1]["repo_id"]

    if selected is None:
        display.info("Cancelled.")
        return None

    return selected


# ── GGUF helpers ──────────────────────────────────────────────────────────────

def _select_gguf_file(
    gguf_files: list[dict],
    quant: str | None,
    no_interactive: bool,
) -> dict:
    """Return the chosen GGUF file entry, or raise SystemExit on failure."""
    if quant:
        quant_upper = quant.upper()
        chosen = next(
            (fe for fe in gguf_files if fe["quant"] == quant_upper),
            None,
        )
        if not chosen:
            available = ", ".join(
                fe["quant"] or "?" for fe in gguf_files
            )
            display.error(
                f"Quantisation [bold]{quant}[/] not found.\n\n"
                f"  Available: {available}"
            )
            raise SystemExit(1)
        return chosen

    if no_interactive:
        for pref in ("Q4_K_M", "Q5_K_M", "Q4_0"):
            chosen = next(
                (fe for fe in gguf_files if fe["quant"] == pref),
                None,
            )
            if chosen:
                return chosen
        return gguf_files[len(gguf_files) // 2]

    # Interactive selection
    display.gguf_file_table(gguf_files)
    display.blank()

    choices = [fe["quant"] or fe["filename"] for fe in gguf_files]
    default_choice = next(
        (c for c in choices if c in ("Q4_K_M", "Q5_K_M")),
        choices[len(choices) // 2],
    )

    try:
        if _IS_WIN:
            raise RuntimeError("skip questionary on Windows")
        import questionary  # type: ignore  # pylint: disable=import-outside-toplevel
        selected = questionary.select(
            "Select quantization:",
            choices=choices,
            default=default_choice,
            style=questionary.Style([
                ("selected", "bold fg:green"),
                ("pointer", "bold fg:green"),
            ]),
        ).ask()
    except Exception:  # pylint: disable=broad-except
        display.info(f"Available: {', '.join(choices)}")
        selected = click.prompt(
            "  Select quantization",
            default=default_choice,
            type=click.Choice(choices, case_sensitive=False),
        )

    if selected is None:
        display.info("Cancelled.")
        raise SystemExit(0)

    return next(
        fe for fe in gguf_files if (fe["quant"] or fe["filename"]) == selected
    )


def _install_gguf(
    info: downloader.ModelInfo,
    repo_id: str,
    quant: str | None,
    dest_dir: Path,
    no_interactive: bool,
) -> None:
    """Download and register a GGUF model."""
    gguf_files = info.gguf_files
    if not gguf_files:
        display.error("No GGUF files found in this repository.")
        raise SystemExit(1)

    chosen = _select_gguf_file(gguf_files, quant, no_interactive)
    filename = chosen["filename"]

    display.blank()
    display.step(f"Downloading [bold]{filename}[/]…")
    start = time.time()

    try:
        local_path = downloader.download_gguf_file(
            repo_id, filename, dest_dir
        )
    except Exception as exc:  # pylint: disable=broad-except
        display.error(_download_error_message(repo_id, exc))
        raise SystemExit(1) from None

    elapsed = time.time() - start
    size_bytes = local_path.stat().st_size

    registry.register(
        model_id=repo_id,
        repo_id=repo_id,
        model_type="gguf",
        path=dest_dir,
        files=[filename],
        primary_file=filename,
        size_bytes=size_bytes,
        tags=info.tags,
        metadata=info.card_data,
    )

    display.blank()
    display.install_summary(repo_id, size_bytes, elapsed)


# ── Transformer install ───────────────────────────────────────────────────────

def _install_transformer(
    info: downloader.ModelInfo,
    repo_id: str,
    dest_dir: Path,
) -> None:
    """Download and register a full Transformer model snapshot."""
    display.step("Downloading model files (this may take a while)…")
    start = time.time()

    try:
        local_dir = downloader.download_transformer_model(repo_id, dest_dir)
    except Exception as exc:  # pylint: disable=broad-except
        display.error(_download_error_message(repo_id, exc))
        raise SystemExit(1) from None

    elapsed = time.time() - start
    size_bytes = sum(
        fp.stat().st_size for fp in local_dir.rglob("*") if fp.is_file()
    )
    file_names = [fp.name for fp in local_dir.iterdir() if fp.is_file()]
    primary = next(
        (fn for fn in file_names if "config.json" in fn),
        file_names[0] if file_names else None,
    )

    registry.register(
        model_id=repo_id,
        repo_id=repo_id,
        model_type="transformers",
        path=local_dir,
        files=file_names,
        primary_file=primary,
        size_bytes=size_bytes,
        tags=info.tags,
        metadata=info.card_data,
    )

    display.blank()
    display.install_summary(repo_id, size_bytes, elapsed)


# ── single-file install ───────────────────────────────────────────────────────

def _install_single_file(
    info: downloader.ModelInfo,
    repo_id: str,
    filename: str,
    dest_dir: Path,
) -> None:
    """Download and register a single named file."""
    display.step(f"Downloading [bold]{filename}[/]…")
    start = time.time()

    try:
        local_path = downloader.download_gguf_file(repo_id, filename, dest_dir)
    except Exception as exc:  # pylint: disable=broad-except
        display.error(_download_error_message(repo_id, exc))
        raise SystemExit(1) from None

    elapsed = time.time() - start
    size_bytes = local_path.stat().st_size
    model_type = "gguf" if filename.lower().endswith(".gguf") else "transformers"

    registry.register(
        model_id=repo_id,
        repo_id=repo_id,
        model_type=model_type,
        path=dest_dir,
        files=[filename],
        primary_file=filename,
        size_bytes=size_bytes,
        tags=info.tags,
        metadata=info.card_data,
    )

    display.blank()
    display.install_summary(repo_id, size_bytes, elapsed)
