"""
Rich-based display helpers — gives llmpm its npm-like aesthetic.
"""

from __future__ import annotations

import math
from typing import Any

from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
from rich import box

console = Console()
err_console = Console(stderr=True)


# ── brand ────────────────────────────────────────────────────────────────────

BRAND = "[bold green]llmpm[/]"


def header(text: str = "") -> None:
    """Print the llmpm brand header with an optional sub-label."""
    console.print()
    if text:
        console.print(f"  {BRAND}  [dim]{text}[/]")
    else:
        console.print(f"  {BRAND}")


def blank() -> None:
    """Print a blank line."""
    console.print()


# ── status lines ─────────────────────────────────────────────────────────────

def step(msg: str) -> None:
    """Print a dim in-progress message."""
    console.print(f"  [dim]{msg}[/]")


def ok(msg: str) -> None:  # pylint: disable=invalid-name
    """Print a green success line."""
    console.print(f"  [green]✓[/]  {msg}")


def warn(msg: str) -> None:
    """Print a yellow warning line."""
    console.print(f"  [yellow]![/]  [yellow]{msg}[/]")


def error(msg: str) -> None:
    """Print a bold red error to stderr."""
    err_console.print(f"\n  [bold red]error[/]  {msg}\n")


def info(msg: str) -> None:
    """Print a dim informational line."""
    console.print(f"  [cyan]i[/]  [dim]{msg}[/]")


# ── model card ───────────────────────────────────────────────────────────────

def backend_badge(model_type: str) -> Text:
    """Return a badge for the serving backend (GGUF, Transformers, Diffusers)."""
    if model_type == "gguf":
        return Text("GGUF", style="bold white on blue")
    return Text("Transformers", style="bold white on magenta")


def pipeline_badge(category: str) -> Text:
    """Return a badge for the detected pipeline category."""
    _badges: dict[str, tuple[str, str]] = {
        "text-generation":     ("Text Gen",  "bold white on magenta"),
        "image-generation":    ("Image Gen", "bold white on dark_orange"),
        "vision":              ("Vision",    "bold white on purple"),
        "audio-transcription": ("ASR",       "bold white on dark_cyan"),
        "audio-speech":        ("TTS",       "bold white on dark_green"),
    }
    label, style = _badges.get(category, ("—", "dim"))
    return Text(label, style=style)


def model_type_badge(model_type: str) -> Text:
    """Legacy single-badge helper (used by llmpm info)."""
    if model_type == "gguf":
        return Text("GGUF", style="bold white on blue")
    return pipeline_badge(model_type)


def model_card(
    *,
    repo_id: str,
    model_type: str,
    tags: list[str],
    total_size: int,
    file_count: int,
    extra: dict[str, Any] | None = None,
) -> None:
    """Print a Rich info card for a model (repo, type, size, tags)."""
    badge = model_type_badge(model_type)

    table = Table.grid(padding=(0, 2))
    table.add_column(style="dim", justify="right")
    table.add_column()

    table.add_row("repo", f"[bold]{repo_id}[/]")
    table.add_row("type", badge)
    plural = "s" if file_count != 1 else ""
    table.add_row("files", f"{file_count} file{plural}")
    if total_size:
        table.add_row("size", _fmt_size(total_size))
    if tags:
        table.add_row("tags", "  ".join(f"[dim]{t}[/]" for t in tags[:6]))
    if extra:
        for key, val in extra.items():
            table.add_row(key, str(val))

    console.print(
        Panel(table, border_style="dim", padding=(0, 2))
    )


# ── GGUF quantisation picker ─────────────────────────────────────────────────

QUANT_DESCRIPTIONS = {
    "Q2_K":   "2-bit  · smallest, lowest quality",
    "Q3_K_S": "3-bit small",
    "Q3_K_M": "3-bit medium",
    "Q3_K_L": "3-bit large",
    "Q4_0":   "4-bit, original",
    "Q4_1":   "4-bit",
    "Q4_K_S": "4-bit small",
    "Q4_K_M": "4-bit medium  · good balance",
    "Q5_0":   "5-bit",
    "Q5_1":   "5-bit",
    "Q5_K_S": "5-bit small",
    "Q5_K_M": "5-bit medium",
    "Q6_K":   "6-bit  · near lossless",
    "Q8_0":   "8-bit  · near lossless, larger",
    "F16":    "16-bit float  · full precision",
    "BF16":   "16-bit bfloat  · full precision",
    "F32":    "32-bit float  · full precision",
}

RECOMMENDED_QUANTS = {"Q4_K_M", "Q5_K_M"}


def gguf_file_table(files: list[dict]) -> None:
    """Print a table of available GGUF quantisations."""
    table = Table(
        show_header=True,
        box=box.SIMPLE,
        padding=(0, 1),
        header_style="bold dim",
    )
    table.add_column("#", style="dim", justify="right", width=3)
    table.add_column("Quantization", min_width=12)
    table.add_column("Size", justify="right", min_width=9)
    table.add_column("Description", style="dim")

    for idx, file_entry in enumerate(files, 1):
        quant = file_entry.get("quant") or "?"
        is_rec = quant in RECOMMENDED_QUANTS
        name_text = Text(quant)
        if is_rec:
            name_text.stylize("bold green")
            name_text.append("  ★ recommended", style="dim green")
        desc = QUANT_DESCRIPTIONS.get(quant, "")
        table.add_row(
            str(idx),
            name_text,
            fmt_size(file_entry["size"]) if file_entry["size"] else "—",
            desc,
        )

    console.print(table)


# ── installed models list ────────────────────────────────────────────────────

def models_table(models: dict) -> None:
    """Print a Rich table of all installed models."""
    if not models:
        console.print("  [dim](no models installed)[/]\n")
        return

    # Lazy import to avoid circular dependency at module load time.
    from pathlib import Path  # pylint: disable=import-outside-toplevel
    from llmpm.core import model_detector  # pylint: disable=import-outside-toplevel

    table = Table(
        show_header=True,
        box=box.ROUNDED,
        border_style="dim",
        padding=(0, 1),
        header_style="bold",
    )
    table.add_column("Model", min_width=30)
    table.add_column("Backend", justify="center", width=12)
    table.add_column("Pipeline", justify="center", width=11)
    table.add_column("File", min_width=20)
    table.add_column("Size", justify="right", min_width=9)
    table.add_column("Installed")

    total_bytes = 0
    for model in models.values():
        raw_type = model.get("model_type", "?")
        backend = backend_badge(raw_type)

        category = raw_type
        if raw_type != "gguf":
            try:
                category = model_detector.detect(Path(model["path"]))
            except Exception:  # pylint: disable=broad-except
                category = raw_type
        pipeline = pipeline_badge(category)

        primary = model.get("primary_file") or "—"
        size = model.get("size_bytes") or 0
        total_bytes += size

        installed_at = model.get("installed_at", "")
        if installed_at:
            installed_at = installed_at[:10]  # date only

        table.add_row(
            f"[bold]{model['repo_id']}[/]",
            backend,
            pipeline,
            f"[dim]{primary}[/]",
            fmt_size(size) if size else "—",
            f"[dim]{installed_at}[/]",
        )

    console.print(table)
    count = len(models)
    console.print(
        f"\n  [dim]{count} model{'s' if count != 1 else ''}"
        f"  ·  {fmt_size(total_bytes)} total[/]\n"
    )


# ── success summary ──────────────────────────────────────────────────────────

def install_summary(model_id: str, size_bytes: int, elapsed: float) -> None:
    """Print the post-install success panel."""
    elapsed_str = fmt_elapsed(elapsed)
    size_str = fmt_size(size_bytes)

    console.print(
        Panel(
            f"  [green]added[/] [bold]{model_id}[/]  "
            f"[dim]({size_str})[/]  [dim]in {elapsed_str}[/]",
            border_style="green",
            padding=(0, 2),
        )
    )
    console.print()


def push_summary(repo_id: str, url: str) -> None:
    """Print the post-push success panel."""
    console.print(
        Panel(
            f"  [green]pushed[/] [bold]{repo_id}[/]\n\n"
            f"  [dim]{url}[/]",
            border_style="green",
            padding=(0, 2),
        )
    )
    console.print()


# ── run header ───────────────────────────────────────────────────────────────

def run_header(
    repo_id: str,
    model_type: str,
    primary_file: str | None,
    n_ctx: int | None = None,
    n_gpu_layers: int | None = None,
    max_tokens: int | None = None,
    system_prompt: str | None = None,
) -> None:
    """Print the model/backend info panel shown before a run/serve session."""
    backend = (
        "llama.cpp" if model_type == "gguf" else "HuggingFace Transformers"
    )
    table = Table.grid(padding=(0, 2))
    table.add_column(style="dim", justify="right")
    table.add_column()

    table.add_row("model", f"[bold]{repo_id}[/]")
    table.add_row("backend", backend)
    if primary_file:
        table.add_row("file", f"[dim]{primary_file}[/]")
    if n_ctx is not None:
        table.add_row("ctx", f"{n_ctx:,} tokens")
    if n_gpu_layers is not None:
        gpu_str = "all" if n_gpu_layers == -1 else str(n_gpu_layers)
        table.add_row("gpu layers", gpu_str)
    if max_tokens is not None:
        table.add_row("max tokens", f"{max_tokens:,}")
    if system_prompt:
        disp = (
            system_prompt if len(system_prompt) <= 60
            else system_prompt[:57] + "…"
        )
        table.add_row("system", f"[dim]{disp}[/]")

    console.print(Panel(table, border_style="dim", padding=(0, 2)))
    console.print()


# ── serve banner ─────────────────────────────────────────────────────────────

def serve_banner(host: str, port: int, category: str = "text-generation") -> None:
    """Print the server address, endpoints, and curl usage examples."""
    curl_host = "localhost" if host in ("0.0.0.0", "") else host
    url = f"http://{curl_host}:{port}"

    blank()
    ok(f"Listening on  {url}")
    blank()
    info("[dim]GET[/]   [bold]/chat[/]                    [green]browser chat UI[/]")
    info("[dim]GET[/]   [bold]/health[/]                 health check")
    info("[dim]GET[/]   [bold]/v1/models[/]              list model")
    info("[dim]GET[/]   [bold]/v1/models/{id}[/]         model info")
    info("[dim]POST[/]  [bold]/v1/chat/completions[/]    chat inference")
    info("[dim]POST[/]  [bold]/v1/completions[/]         legacy text completion")
    info("[dim]POST[/]  [bold]/v1/embeddings[/]          text embeddings")
    info("[dim]POST[/]  [bold]/v1/images/generations[/]  text-to-image")
    info("[dim]POST[/]  [bold]/v1/audio/transcriptions[/] speech-to-text")
    info("[dim]POST[/]  [bold]/v1/audio/speech[/]        text-to-speech")
    blank()
    step("Examples:")
    blank()
    step(f"  open {url}/chat")
    blank()
    step(f"  curl {url}/health")
    blank()

    if category == "image-generation":
        step(
            f"  curl -X POST {url}/v1/images/generations \\\n"
            "    -H \"Content-Type: application/json\" \\\n"
            "    -d '{\"prompt\": \"summer in the mountains\"}'\n"
            "\n"
            "  # Response:\n"
            "  # {\n"
            "  #   \"created\": 1234567890,\n"
            "  #   \"data\": [{\"b64_json\": \"<base64 PNG>\"}]\n"
            "  # }"
        )
    elif category == "audio-transcription":
        step(
            f"  curl -X POST {url}/v1/audio/transcriptions \\\n"
            "    -H \"Content-Type: application/octet-stream\" \\\n"
            "    --data-binary @audio.wav\n"
            "\n"
            "  # Response:\n"
            "  # {\"text\": \"transcribed text here\"}"
        )
    elif category == "audio-speech":
        step(
            f"  curl -X POST {url}/v1/audio/speech \\\n"
            "    -H \"Content-Type: application/json\" \\\n"
            "    -d '{\"input\": \"Hello world\"}' \\\n"
            "    --output speech.wav"
        )
    else:
        step(
            f"  curl -X POST {url}/v1/chat/completions \\\n"
            "    -H \"Content-Type: application/json\" \\\n"
            "    -d '{\"messages\": "
            "[{\"role\": \"user\", \"content\": \"Hello!\"}]}'"
        )

    blank()
    step("Press Ctrl-C to stop.")
    blank()


# ── utilities ────────────────────────────────────────────────────────────────

def fmt_size(num_bytes: int) -> str:
    """Format a byte count as a human-readable string (e.g. 1.24 GB)."""
    if num_bytes == 0:
        return "0 B"
    units = ["B", "KB", "MB", "GB", "TB"]
    idx = min(int(math.log(num_bytes, 1024)), len(units) - 1)
    value = num_bytes / (1024 ** idx)
    if value >= 100:
        return f"{value:.0f} {units[idx]}"
    if value >= 10:
        return f"{value:.1f} {units[idx]}"
    return f"{value:.2f} {units[idx]}"


def fmt_elapsed(seconds_total: float) -> str:
    """Format a duration in seconds as a human-readable string."""
    if seconds_total < 60:
        return f"{seconds_total:.0f}s"
    mins = int(seconds_total // 60)
    secs = int(seconds_total % 60)
    if mins < 60:
        return f"{mins}m {secs}s"
    hours = mins // 60
    mins = mins % 60
    return f"{hours}h {mins}m"


# Keep private alias for backward compatibility.
_fmt_size = fmt_size
