"""
llmpm serve <repo_id>

Starts a local HTTP server exposing the model as a REST API.
Compatible with the OpenAI /v1/chat/completions format.

Examples:
  llmpm serve bartowski/Llama-3.2-1B-Instruct-GGUF
  llmpm serve bartowski/Llama-3.2-1B-Instruct-GGUF --port 9000
  llmpm serve meta-llama/Llama-3.2-1B-Instruct --host 127.0.0.1
"""

from __future__ import annotations

import socket
from pathlib import Path

import click

from llmpm import display
from llmpm.core import model_detector, registry
from llmpm.core import (
    serve_audio,
    serve_diffusion,
    serve_gguf,
    serve_transformers,
    serve_vision,
)


@click.command("serve")
@click.argument("repo_id")
@click.option(
    "--port", "-p",
    default=8080,
    show_default=True,
    help="Port to listen on.",
)
@click.option(
    "--host", "-H",
    default="localhost",
    show_default=True,
    help="Host/address to bind to.",
)
@click.option(
    "--ctx",
    default=128000,
    show_default=True,
    help="Context window size (GGUF only).",
)
@click.option(
    "--gpu-layers", "gpu_layers",
    default=-1,
    show_default=True,
    help="GPU layers to offload (GGUF only, -1 = all).",
)
@click.option(
    "--max-tokens", "max_tokens",
    default=128000,
    show_default=True,
    help="Default max tokens per response (overridable per-request).",
)
def serve(
    repo_id: str,
    port: int,
    host: str,
    ctx: int,
    gpu_layers: int,
    max_tokens: int,
) -> None:
    """Serve a model over HTTP (OpenAI-compatible API)."""
    display.header(f"serve {repo_id}")
    display.blank()

    match = registry.find_model(repo_id)
    if not match:
        display.error(
            f"No installed model matches [bold]{repo_id}[/].\n\n"
            f"  Install first:  [bold cyan]llmpm install {repo_id}[/]"
        )
        raise SystemExit(1)

    matched_id, entry = match
    if matched_id != repo_id:
        display.info(f"Using [bold]{matched_id}[/]")
    repo_id = matched_id

    model_type = entry["model_type"]
    primary_file = entry.get("primary_file")
    model_path = Path(entry["path"])

    display.run_header(
        repo_id=repo_id,
        model_type=model_type,
        primary_file=primary_file,
        n_ctx=ctx if model_type == "gguf" else None,
        n_gpu_layers=gpu_layers if model_type == "gguf" else None,
        max_tokens=max_tokens,
    )

    actual_port = _find_free_port(host, port)
    if actual_port != port:
        display.warn(f"Port {port} is busy — using {actual_port} instead.")

    if model_type == "gguf":
        gguf_path = model_path / primary_file if primary_file else model_path
        serve_gguf.serve(
            repo_id, gguf_path, host, actual_port, ctx, gpu_layers, max_tokens,
        )
        return

    # ── HuggingFace model: detect pipeline category ──────────────────────────
    category = model_detector.detect(model_path)
    display.info(
        f"Detected: [bold]{model_detector.label(category)}[/]"
    )

    if category == model_detector.IMAGE_GENERATION:
        serve_diffusion.serve(
            repo_id, model_path, host, actual_port, max_tokens,
        )
    elif category == model_detector.VISION:
        serve_vision.serve(
            repo_id, model_path, host, actual_port, max_tokens,
        )
    elif category == model_detector.AUDIO_TRANSCRIPTION:
        serve_audio.serve(
            repo_id, model_path, host, actual_port, max_tokens,
            pipeline_tag="automatic-speech-recognition",
        )
    elif category == model_detector.AUDIO_SPEECH:
        serve_audio.serve(
            repo_id, model_path, host, actual_port, max_tokens,
            pipeline_tag="text-to-speech",
        )
    else:
        try:
            serve_transformers.serve(
                repo_id, model_path, host, actual_port, max_tokens,
            )
        except serve_transformers.ModelNotTextGenerationError:
            # Detection was wrong — re-detect by forcing a fresh scan and
            # trying the diffusion backend as the most likely alternative.
            display.warn(
                "Text-generation load failed; attempting diffusion backend…"
            )
            serve_diffusion.serve(
                repo_id, model_path, host, actual_port, max_tokens,
            )


# ── helpers ──────────────────────────────────────────────────────────────────

def _find_free_port(host: str, preferred: int) -> int:
    """Return preferred port if available, else the next free one."""
    port = preferred
    while port < 65535:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            try:
                sock.bind((host, port))
                return port
            except OSError:
                port += 1
    return preferred
