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

from . 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, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> 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
            continuation_text: Part of the response that is already generated
            should_stop: Callback that returns True when inference should stop

        Returns:
            Iterator[str]: An iterator that yields the generated text.
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context},
            {"role": "assistant", "content": continuation_text},
        ]
        
        stream = self._client.chat_completion(
            model=self._model,
            messages=messages,
            temperature=self._temperature,
            add_generation_prompt=False,
            stream=True
        )
        
        try:
            for response in stream:
                if should_stop():
                    stream.close()
                    break
                if content := response.choices[0].delta.content:
                    yield content
        finally:
            stream.close()

    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