from typing import Iterator, Optional, Callable
from mistralai import Mistral
from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer

from . import LlmEngine
from ..util import skip_prefix

class MistralLlmEngine(LlmEngine):
    def __init__(
            self,
            model: str,
            temperature: float,
            token_limit: int,
            api_key: str,
        ):
        self._model = model
        self._temperature = temperature
        self._token_limit = token_limit
        self._api_key = api_key
        self._client = Mistral(api_key=api_key)
        self._tokenizer = MistralTokenizer.v3()

    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
        messages = [
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": main_context,
            },
            {
                "role": "assistant",
                "content": continuation_text,
                "prefix": True,
            },
        ] if continuation_text else [
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": main_context,
            },
        ]
        stream_response = self._client.chat.stream(
            model=self._model,
            messages=messages,
            temperature=self._temperature,
        )
        
        try:
            def content_generator():
                for chunk in stream_response:
                    if should_stop():
                        stream_response.response.close()
                        break
                    if content := chunk.data.choices[0].delta.content:
                        yield content
            yield from skip_prefix(content_generator(), continuation_text)
        finally:
            stream_response.response.close()

    def token_count(self, system_prompt: str, main_context: str) -> int:
        messages = [
            SystemMessage(content=system_prompt),
            UserMessage(content=main_context),
        ]

        tokenized = self._tokenizer.encode_chat_completion(
            ChatCompletionRequest(
                messages=messages,
                model=self._model
            )
        )
        return len(tokenized.tokens)

    def token_limit(self) -> int:
        return self._token_limit