"""
llmpm serve [model1] [model2] ...

Starts a single local HTTP server exposing one or more models as a REST API.
Compatible with the OpenAI /v1/chat/completions format.

When multiple models are loaded the caller selects the active model via the
``model`` field in the request body; if omitted the first loaded model is used.

If no model names are given every installed model is served.

Examples:
  llmpm serve bartowski/Llama-3.2-1B-Instruct-GGUF
  llmpm serve model-a model-b --port 9000
  llmpm serve                             # serve ALL installed models
  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.commands._picker import pick_installed_model
from llmpm.core import model_detector, registry
from llmpm.core.downloader import normalize_repo_id
from llmpm.core import (
    serve_audio,
    serve_diffusion,
    serve_gguf,
    serve_transformers,
    serve_vision,
)
from llmpm.core import server as _server


@click.command("serve")
@click.argument("repo_ids", nargs=-1)
@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).",
)
@click.option(
    "--path",
    "extra_paths",
    multiple=True,
    type=click.Path(exists=True, file_okay=True, dir_okay=True),
    help="Path to a local model directory or GGUF file (bypasses registry). "
         "Can be repeated to serve multiple custom-path models.",
)
def serve(
    repo_ids: tuple[str, ...],
    port: int,
    host: str,
    ctx: int,
    gpu_layers: int,
    max_tokens: int,
    extra_paths: tuple[str, ...],
) -> None:
    """Serve one or more models over HTTP (OpenAI-compatible API).

    If no model names are given, every installed model is served on a single
    server.  Clients select a model via the ``model`` field in the request
    body.

    Use --path to serve a model that is not in the llmpm registry:

      llmpm serve --path /path/to/model
      llmpm serve --path /path/to/model.gguf
    """
    # ── resolve model list ───────────────────────────────────────────────────
    has_registry_ids = bool(repo_ids)
    has_paths = bool(extra_paths)

    if not has_registry_ids and not has_paths:
        all_models = registry.all_models()
        if not all_models:
            display.error(
                "No models installed.\n\n"
                "  Install one first:  [bold cyan]llmpm install <model>[/]\n"
                "  Or serve a local path:  "
                "[bold cyan]llmpm serve --path /path/to/model[/]"
            )
            raise SystemExit(1)
        resolved_ids = list(all_models.keys())
        display.header("serve  [dim](all installed models)[/]")
    else:
        label_parts = list(repo_ids) + [Path(p).name for p in extra_paths]
        display.header(f"serve  [dim]{' · '.join(label_parts)}[/]")
        resolved_ids = [normalize_repo_id(r) for r in repo_ids]

    display.blank()

    # ── fuzzy-match + pick each requested model ──────────────────────────────
    entries: list[tuple[str, dict]] = []
    for query in resolved_ids:
        matches = registry.find_models(query)
        if not matches:
            display.error(
                f"No installed model matches [bold]{query}[/].\n\n"
                f"  Install first:  [bold cyan]llmpm install {query}[/]"
            )
            raise SystemExit(1)

        if len(matches) == 1:
            matched_id, entry = matches[0]
            if matched_id != query:
                display.info(f"Using [bold]{matched_id}[/]")
        else:
            result = pick_installed_model(matches, query)
            if result is None:
                raise SystemExit(0)
            matched_id, entry = result

        entries.append((matched_id, entry))

    # ── add entries from --path arguments ────────────────────────────────────
    for raw_path in extra_paths:
        p = Path(raw_path).expanduser().resolve()
        if p.is_file() and p.suffix.lower() == ".gguf":
            model_id = p.stem
            entry = {
                "model_type": "gguf",
                "path": str(p.parent),
                "primary_file": p.name,
            }
        elif p.is_dir():
            model_id = p.name
            gguf_files = list(p.glob("*.gguf"))
            entry = {
                "model_type": "gguf" if gguf_files else "transformers",
                "path": str(p),
                "primary_file": None,
            }
        else:
            display.error(
                f"[bold]{raw_path}[/] is not a directory or .gguf file."
            )
            raise SystemExit(1)
        display.info(
            f"Using local path: [bold]{p}[/]  [dim](not from registry)[/]"
        )
        entries.append((model_id, entry))

    # ── announce what will be loaded (before weights are pulled into RAM) ────
    if len(entries) > 1:
        display.info(
            f"Starting [bold]{len(entries)}[/] models:"
        )
        for rid, ent in entries:
            mtype = ent.get("model_type", "transformer")
            display.info(f"  · [bold]{rid}[/]  [dim]{mtype}[/]")
        display.blank()

    # ── load each model ──────────────────────────────────────────────────────
    contexts: list[_server.HandlerContext] = []

    for repo_id, entry in entries:
        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,
        )

        if not model_path.exists():
            display.error(
                f"Model [bold]{repo_id}[/] is registered but its files are missing "
                f"at [bold]{model_path}[/].\n\n"
                f"  Re-install with:  [bold cyan]llmpm install {repo_id}[/]"
            )
            raise SystemExit(1)

        if model_type == "gguf":
            gguf_path = (
                model_path / primary_file if primary_file else model_path
            )
            context = serve_gguf.load(
                repo_id, gguf_path, ctx, gpu_layers, max_tokens,
                metadata=dict(entry),
            )
            contexts.append(context)
            continue

        # ── 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:
            context = serve_diffusion.load(
                repo_id, model_path, max_tokens, metadata=dict(entry),
            )
        elif category == model_detector.VISION:
            context = serve_vision.load(
                repo_id, model_path, max_tokens, metadata=dict(entry),
            )
        elif category == model_detector.AUDIO_TRANSCRIPTION:
            context = serve_audio.load(
                repo_id, model_path, max_tokens,
                pipeline_tag="automatic-speech-recognition",
                metadata=dict(entry),
            )
        elif category == model_detector.AUDIO_SPEECH:
            context = serve_audio.load(
                repo_id, model_path, max_tokens,
                pipeline_tag="text-to-speech",
                metadata=dict(entry),
            )
        else:
            try:
                context = serve_transformers.load(
                    repo_id, model_path, max_tokens,
                    metadata=dict(entry),
                )
            except serve_transformers.ModelNotTextGenerationError:
                display.warn(
                    "Text-generation load failed; "
                    "attempting diffusion backend…"
                )
                context = serve_diffusion.load(
                    repo_id, model_path, max_tokens,
                    metadata=dict(entry),
                )

        contexts.append(context)

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

    first_category = contexts[0].category if contexts else "text-generation"
    display.serve_banner(host, actual_port, category=first_category)

    if len(contexts) > 1:
        display.info(
            f"Serving [bold]{len(contexts)}[/] models — "
            "select via [bold]model[/] field in request body"
        )
        for c in contexts:
            display.info(f"  · [bold]{c.model_id}[/]  [dim]{c.category}[/]")
        display.blank()

    _server.start_multi(contexts, host=host, port=actual_port)


# ── 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
