from typing import Iterator, Optional
from huggingface_hub import InferenceClient

from .llm_engine import LlmEngine

class HfLlmEngine(LlmEngine):
    """
    LLM Engine implementation using HuggingFace's InferenceClient.
    """
    
    def __init__(
        self,
        model_id: str = "mistralai/Mistral-7B-Instruct-v0.2",
        api_token: Optional[str] = None,
        temperature: float = 0.7,
        max_new_tokens: int = 1024,
    ):
        """
        Initialize the HuggingFace Inference API LLM Engine.

        Args:
            model_id: HuggingFace model ID to use (default: Mistral-7B-Instruct)
            api_token: HuggingFace API token. If None, will try to read from HF_TOKEN env var
            temperature: Sampling temperature (default: 0.7)
            max_new_tokens: Maximum number of tokens to generate (default: 1024)
        """
        self.model_id = model_id
        self.client = InferenceClient(token=api_token)
        
        # Generation parameters
        self.temperature = temperature
        self.max_new_tokens = max_new_tokens
            
    def set_model_path(self, model_id: str):
        """
        Update the model being used.

        Args:
            model_id: New HuggingFace model ID to use
        """
        self.model_id = model_id
                        
    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.
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": main_context}
        ]
        
        def stream_wrapper():
            stream = self.client.chat_completion(
                model=self.model_id,
                messages=messages,
                temperature=self.temperature,
                max_tokens=self.max_new_tokens,
                stream=True
            )
            
            for response in stream:
                if content := response.choices[0].delta.content:
                    yield content
        return stream_wrapper()