from typing import Iterator, Optional
from abc import ABC, abstractmethod
import os
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 .llm_engine import LlmEngine

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) -> Iterator[str]:
        messages = [
            SystemMessage(content=system_prompt),
            UserMessage(content=main_context),
        ]
        messages = [
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": main_context,
            },
            {
                "role": "assistant",
                "content": "<",
                "prefix": True,
            },
        ]
        stream_response = self._client.chat.stream(
            model=self._model,
            messages=messages,
            temperature=self._temperature,
        )
        for chunk in stream_response:
            if content := chunk.data.choices[0].delta.content:
                yield chunk.data.choices[0].delta.content

    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