"""
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 time
from pathlib import Path

import click

from llmpm import display
from llmpm.core import downloader, registry


@click.command("install")
@click.argument("repo_id")
@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.",
)
def install(
    repo_id: str,
    quant: str | None,
    specific_file: str | None,
    no_interactive: bool,
) -> None:
    """Download and install a model from HuggingFace Hub."""
    display.header(f"install {repo_id}")
    display.blank()

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

    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

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


# ── 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:
        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 ImportError:
        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(f"Download failed: {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(f"Download failed: {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(f"Download failed: {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)
