from typing import Callable, Iterator
import openai
import tiktoken

from .llm_engine import LlmEngine

class OpenAILlmEngine(LlmEngine):
    """
    LLM Engine implementation using OpenAI's API.
    Supports streaming responses from chat completion models.
    """
    
    def __init__(
        self,
        model: str,
        temperature: float,
        token_limit: int,
        api_key: str,
    ):
        """
        Initialize the OpenAI LLM Engine.

        Args:
            model: OpenAI model to use
            temperature: Temperature for sampling
            api_key: OpenAI API key
            token_limit: Maximum number of tokens to generate
        """
        self._model = model
        self._temperature = temperature
        self._token_limit = token_limit

        self._client = openai.Client(
            api_key=api_key,
        )
                        
    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context}
        ]
        
        stream = self._client.chat.completions.create(
            model=self._model,
            messages=messages,
            temperature=self._temperature,
            stream=True,
        )

        try:
            for chunk in stream:
                if should_stop():
                    break
                if content := chunk.choices[0].delta.content:
                    yield content
        finally:
            stream.close()
            #stream.response.close()
    
    def token_count(self, system_prompt: str, main_context: str) -> int:
        """
        Calculate the total token count for the system prompt and context.
        
        Args:
            system_prompt: The system prompt string
            main_context: The main context string
            
        Returns:
            int: Total number of tokens
        """
        encoding = tiktoken.encoding_for_model(self._model)
        return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
    
    def token_limit(self) -> int:
        return self._token_limit