from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoConfig
from typing import Iterator, Optional

from .llm_engine import LlmEngine

class HfLlmEngine(LlmEngine):
    """
    LLM Engine implementation using HuggingFace's InferenceClient.
    """
    
    def __init__(
        self,
        model: str,
        temperature: float,
        api_token: Optional[str],
    ):
        """
        Initialize the HuggingFace Inference API LLM Engine.

        Args:
            model: HuggingFace model ID to use
            temperature: Sampling temperature
            api_token: HuggingFace API token
        """
        self._model = model
        self._temperature = temperature

        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
        self._config = AutoConfig.from_pretrained(model, token=api_token)
        self._client = InferenceClient(token=api_token)
                        
    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
        """
        Run inference using the system prompt and main context.

        Args:
            system_prompt: The system prompt string
            main_context: The main context string after templating

        Returns:
            Iterator[str]: An iterator that yields the generated text.
        """
        token_count=self.token_count(system_prompt, main_context)

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context}
        ]
        
        def stream_wrapper():
            stream = self._client.chat_completion(
                model=self._model,
                messages=messages,
                temperature=self._temperature,
                stream=True
            )
            
            for response in stream:
                if content := response.choices[0].delta.content:
                    yield content
        return stream_wrapper()

    def token_count(self, system_prompt: str, main_context: str) -> int:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context}
        ]
        prompt = self._tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        return len(self._tokenizer.encode(prompt))

    def token_limit(self) -> int:
        """
        Get the model's context window size.
        
        Returns:
            int: Maximum number of tokens the model can process
        """
        return self._config.max_position_embeddings