"""llama.cpp runtime wrappers -- the only module that imports llama_cpp.

Both runtimes expose the same `chat(messages, grammar) -> str` contract the
conversation engine depends on, and let the model's own chat template (from the
GGUF metadata) format the messages, so swapping Nemotron <-> MiniCPM is a config
change.

- `LlamaRuntime`: loads the model once at startup. For local / CPU / a plain GPU
  Space where the model can persist in memory across turns.
- `ZeroGpuLlamaRuntime`: for HF Spaces ZeroGPU, where a real GPU exists only inside
  an `@spaces.GPU` function. The model is therefore constructed *inside* that
  function on every call (the llama.cpp maintainer's documented pattern). This is
  cheap to do correctly because the engine holds no conversation state in the
  runtime -- `conversation.respond` rebuilds the full prompt each turn.
"""

from __future__ import annotations

import os

# The model's `<state>` block ends a generation; we stop on it for speed, then
# restore the tag so the parser regex still matches. Shared by both runtimes.
STOP = ["</state>"]


def _finalize(text: str) -> str:
    if "<state>" in text and "</state>" not in text:
        text += "</state>"
    return text


def _complete(llm, messages: list[dict], grammar: str, temperature: float, max_tokens: int) -> str:
    """Grammar-constrained completion shared by both runtimes; returns finalized text."""
    from llama_cpp import LlamaGrammar

    compiled = LlamaGrammar.from_string(grammar, verbose=False)
    result = llm.create_chat_completion(
        messages=messages,
        grammar=compiled,
        temperature=temperature,
        max_tokens=max_tokens,
        stop=STOP,
    )
    return _finalize(result["choices"][0]["message"]["content"])


# --- ZeroGPU entrypoint -----------------------------------------------------
# ZeroGPU requires at least one @spaces.GPU function to be DECORATED AT STARTUP
# (i.e. at import time). Decorating lazily on the first request fails with
# "No @spaces.GPU function detected during startup", and a function decorated after
# startup never actually receives a GPU. So the GPU entrypoint lives at module scope
# and is decorated when this module is imported -- app.py imports the module eagerly
# on a Space (see `_on_zerogpu`) to make that happen before ZeroGPU's startup scan.
#
# `spaces` is imported unconditionally: off a ZeroGPU Space it is a true no-op
# (@spaces.GPU returns the function unchanged), so no try/except shim is needed --
# we just declare `spaces` as a dependency so it's installed in local dev too.
import spaces


def _gpu_duration() -> int:
    try:
        return int(os.environ.get("SECOND_DEGREE_GPU_DURATION", "120"))
    except ValueError:
        return 120


@spaces.GPU(duration=_gpu_duration())
def _zerogpu_chat(
    model_path: str,
    messages: list[dict],
    grammar: str,
    n_ctx: int,
    temperature: float,
    max_tokens: int,
    flash_attn: bool,
) -> str:
    """Build the model on the GPU attached for this call and generate one turn.

    A real GPU exists only inside this @spaces.GPU function, so the `Llama` object is
    constructed per call with all layers offloaded -- the llama.cpp maintainer's
    documented ZeroGPU pattern. Correct because `conversation.respond` rebuilds the
    full prompt each turn and keeps no conversation state in the runtime.
    """
    from llama_cpp import Llama

    llm = Llama(
        model_path=model_path,
        n_gpu_layers=-1,      # offload everything to the attached GPU
        n_ctx=n_ctx,
        flash_attn=flash_attn,
        verbose=False,
    )
    return _complete(llm, messages, grammar, temperature, max_tokens)


class LlamaRuntime:
    def __init__(
        self,
        model_path: str,
        n_ctx: int = 4096,
        n_gpu_layers: int = 0,
        temperature: float = 0.8,
        max_tokens: int = 220,
        seed: int | None = None,
        verbose: bool = False,
    ):
        from llama_cpp import Llama  # imported here so the core engine stays dep-free

        self.temperature = temperature
        self.max_tokens = max_tokens
        self._llm = Llama(
            model_path=model_path,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            seed=seed if seed is not None else -1,
            verbose=verbose,
        )

    def chat(self, messages: list[dict], grammar: str) -> str:
        """Generate one grammar-constrained NPC turn and return the raw text."""
        return _complete(self._llm, messages, grammar, self.temperature, self.max_tokens)


class ZeroGpuLlamaRuntime:
    """llama.cpp on HF Spaces ZeroGPU: a thin wrapper over the module-level
    `@spaces.GPU` entrypoint `_zerogpu_chat`.

    The decorated function must be registered at startup (see the note above
    `_zerogpu_chat`), so it lives at module scope rather than being built here.
    This class just holds the per-Space config and forwards each turn to it.

    A real GPU is attached only for the duration of the decorated call, so the
    `Llama` object is constructed per turn with all layers offloaded. Slower per
    turn (model reload), but correct -- and identical in output to the persistent
    runtime, since the full prompt is rebuilt each turn upstream.
    """

    # The @spaces.GPU duration is set at decoration time from SECOND_DEGREE_GPU_DURATION
    # (see _gpu_duration / app.py), not here -- this class only forwards per-turn config.
    def __init__(
        self,
        model_path: str,
        n_ctx: int = 4096,
        temperature: float = 0.8,
        max_tokens: int = 220,
        flash_attn: bool = True,
    ):
        self.model_path = model_path
        self.n_ctx = n_ctx
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.flash_attn = flash_attn

    def chat(self, messages: list[dict], grammar: str) -> str:
        return _zerogpu_chat(
            self.model_path,
            messages,
            grammar,
            self.n_ctx,
            self.temperature,
            self.max_tokens,
            self.flash_attn,
        )