"""LLM client: thin wrapper around llama-cpp-python (local Qwen2.5).

Loads the GGUF once and exposes chat(messages, tools) -> raw model text.
Tools are rendered into the prompt by the model's OWN chat template (so the
format matches what Qwen was trained on); we parse the <tool_call> output
ourselves in the next step, because llama-cpp's Qwen tool-call parsing is
unreliable -- it sometimes leaves the tag in the raw content instead of
filling message["tool_calls"].
"""
from __future__ import annotations

import json
import os
import re
from collections.abc import Mapping, Sequence
from typing import Any, cast

from llama_cpp import CreateChatCompletionResponse, Llama, LlamaGrammar

from config import MODEL_FILE, MODEL_PATH, MODEL_REPO, N_CTX, N_GPU_LAYERS, NARRATOR_TEMP


def _resolve_model(path: str | None) -> str:
    """Local GGUF if present (dev), else pull it from the Hub repo and cache (Space)."""
    path = path or MODEL_PATH
    if os.path.exists(path):
        return path
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)


class LLM:
    def __init__(self, model_path: str | None = None, n_ctx: int = N_CTX):
        self.llm = Llama(
            model_path=_resolve_model(model_path),
            n_gpu_layers=N_GPU_LAYERS,
            n_ctx=n_ctx,
            verbose=False,
        )

    def chat(self, messages: Sequence[Mapping[str, Any]],
             tools: Sequence[Mapping[str, Any]] | None = None,
             temperature: float = NARRATOR_TEMP, max_tokens: int = 512,
             grammar: str | None = None) -> str:
        """Return the model's raw text reply (may contain <tool_call> tags).

        Public API takes plain message/tool dicts (and a GBNF grammar as plain
        text); we cast once here at the llama-cpp boundary (its stubs use stricter
        TypedDict unions) and build the provider-specific LlamaGrammar here too, so
        callers never import llama-cpp.
        """
        out = cast(
            CreateChatCompletionResponse,
            self.llm.create_chat_completion(
                messages=cast(Any, messages),
                tools=cast(Any, tools),
                temperature=temperature,
                max_tokens=max_tokens,
                grammar=LlamaGrammar.from_string(grammar) if grammar else None,
            ),
        )
        return out["choices"][0]["message"]["content"] or ""


_TOOL_CALL_RE = re.compile(r"<tool_call>\s*(.*?)\s*</tool_call>", re.DOTALL)


def parse_tool_calls(text: str) -> list[dict]:
    """Extract Qwen <tool_call>{...}</tool_call> blocks -> [{name, arguments}, ...].

    Relies on the closing tag (not brace matching) so nested argument objects
    parse correctly. Malformed JSON inside a block is skipped.
    """
    calls = []
    for body in _TOOL_CALL_RE.findall(text):
        try:
            calls.append(json.loads(body))
        except json.JSONDecodeError:
            continue
    return calls