import json
import os
from collections.abc import Iterator
from contextlib import redirect_stderr
from functools import wraps
from pathlib import Path

from llama_cpp import Llama

from yera.config2.dataclasses import LLMConfig
from yera.config2.keyring import DevKeyring
from yera.models.data_classes import Message
from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel


def assert_started(fn):
    @wraps(fn)
    def wrapper(self, *args, **kwargs):
        if self._model is None:
            raise ValueError("Model not started. Call model.start() first.")
        return fn(self, *args, **kwargs)

    return wrapper


class LlamaCppLLM(BaseLLMInterface):
    def __init__(self, config: LLMConfig, overrides, creds_map):
        # make kwargs from config and overrides
        kws = {**config.inference.model_dump(), **overrides}
        super().__init__(**kws)
        self._model = None
        models_dir = Path(DevKeyring.get("providers." + creds_map["models_dir"]))

        with (models_dir / "catalogue.json").open() as f:
            catalogue = json.load(f)
        self.model_path = Path(catalogue[config.id])
        if not self.model_path.exists():
            raise FileNotFoundError(self.model_path)

    def start(self):
        if self._model:
            raise ValueError("Model already started.")

        with open(os.devnull, "w") as devnull, redirect_stderr(devnull):
            self._model = Llama(
                model_path=str(self.model_path),
                n_gpu_layers=self.kwargs.get("n_gpu_layers", -1),
                n_ctx=self.kwargs.get("context_length", 8096),
                verbose=False,
            )

    @assert_started
    def stop(self):
        self._model.close()

    @staticmethod
    def _process_stream(stream):
        stream = (x["choices"][0]["delta"] for x in stream)
        stream = (x.get("content") for x in stream)
        stream = (x for x in stream if x is not None)
        return stream

    @assert_started
    def chat(
        self,
        messages: list[Message],
        *,
        temperature: float = 0.2,
        top_p: float = 0.95,
        top_k: int = 40,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        stop: str | list[str] | None = None,
        seed: int | None = None,
        max_tokens: int | None = None,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0,
        repeat_penalty: float = 1.0,
        **llama_cpp_kw,
    ) -> Iterator[str]:
        stream = self._model.create_chat_completion(
            messages,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            typical_p=typical_p,
            stop=stop or [],
            seed=seed,
            max_tokens=max_tokens,
            presence_penalty=presence_penalty,
            frequency_penalty=frequency_penalty,
            repeat_penalty=repeat_penalty,
            stream=True,
            **llama_cpp_kw,
        )
        return self._process_stream(stream)

    @assert_started
    def make_struct(
        self,
        messages: list[Message],
        cls: type[TBaseModel],
        *,
        temperature: float = 0.2,
        top_p: float = 0.95,
        top_k: int = 40,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        seed: int | None = None,
        max_tokens: int | None = None,
        presence_penalty: float = 0.0,
        frequency_penalty: float = 0.0,
        repeat_penalty: float = 1.0,
        **llama_cpp_kw,
    ) -> Iterator[str]:
        stream = self._model.create_chat_completion(
            messages=messages,
            response_format={
                "type": "json_object",
                "schema": cls.model_json_schema(),
            },
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            typical_p=typical_p,
            seed=seed,
            max_tokens=max_tokens,
            presence_penalty=presence_penalty,
            frequency_penalty=frequency_penalty,
            repeat_penalty=repeat_penalty,
            stream=True,
            **llama_cpp_kw,
        )
        return self._process_stream(stream)
