"""
GGUF model serving via llama-cpp-python.
"""

from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path

from llmpm import display
from llmpm.core import runner
from llmpm.core import server as _server


def load(
    repo_id: str,
    model_path: Path,
    n_ctx: int,
    n_gpu_layers: int,
    default_max_tokens: int,
    metadata: dict | None = None,
) -> _server.HandlerContext:
    """Load a GGUF model and return a HandlerContext (does not start server)."""
    LlamaClass = runner.load_llama_cpp()  # pylint: disable=invalid-name

    display.step(f"Loading {repo_id}…")
    llm = LlamaClass(
        model_path=str(model_path),
        n_ctx=n_ctx,
        n_gpu_layers=n_gpu_layers,
        verbose=False,
    )
    display.ok(f"Model ready  [dim]{repo_id}[/]")

    def infer(
        messages: list[dict],
        max_tokens: int,
        temperature: float,
    ) -> str:
        result = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
        )
        return result["choices"][0]["message"]["content"]

    def infer_stream(
        messages: list[dict],
        max_tokens: int,
        temperature: float,
    ) -> Iterator[str]:
        stream = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
        )
        for chunk in stream:
            token = chunk["choices"][0]["delta"].get("content", "")
            if token:
                yield token

    def embed(texts: list[str]) -> list[list[float]]:
        result = llm.create_embedding(texts)
        return [item["embedding"] for item in result["data"]]

    return _server.HandlerContext(
        model_id=repo_id,
        default_max_tokens=default_max_tokens,
        category="text-generation",
        infer=infer,
        infer_stream=infer_stream,
        embed=embed,
        metadata=metadata or {},
    )


def serve(  # pylint: disable=too-many-arguments
    repo_id: str,
    model_path: Path,
    host: str,
    port: int,
    n_ctx: int,
    n_gpu_layers: int,
    default_max_tokens: int,
) -> None:
    """Load a GGUF model and start the HTTP server (single-model convenience)."""
    ctx = load(repo_id, model_path, n_ctx, n_gpu_layers, default_max_tokens)
    display.serve_banner(host, port)
    _server.start_multi([ctx], host=host, port=port)
