"""
Hugging Face AI endpoint implementation.

This module provides integration with Hugging Face's Inference API for LLM inference.
"""

from huggingface_hub import InferenceClient
from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError

DEFAULT_MODEL = "meta-llama/Llama-3.2-3B-Instruct"

class HuggingfaceEndpoint(BaseAIEndpoint):
    """Hugging Face endpoint for cloud-based LLM inference."""

    def _initialize_client(self) -> None:
        """Initialize the Hugging Face client."""
        api_key = self.ai_config.get("api_key", "")
        if not api_key:
            raise AIEndpointRequestError("Hugging Face API key is required")

        # Default timeout of 30 seconds, configurable via ai_config
        timeout = self.ai_config.get("timeout", 30)
        self.client = InferenceClient(
            model=self.model,
            token=api_key,
            timeout=timeout
        )

    def _get_default_model(self) -> str:
        """Get the default Hugging Face model."""
        return DEFAULT_MODEL

    def query(self, prompt: str, output_format: dict) -> str:
        """
        Send a query to Hugging Face and return the response.

        Args:
            prompt: The prompt to send to the model

        Returns:
            The model's response as a string

        Raises:
            AIEndpointRequestError: If the request fails
        """
        try:
            response = self.client.chat_completion(
                messages=[{"role": "user", "content": prompt}],
                max_tokens=self.max_tokens,
                temperature=self.temperature,
                response_format= {
                    "type": "json_schema",
                    "json_schema": {
                        "name": "output_format",
                        "schema": output_format.model_json_schema(),
                        "strict": True,
                    }
                }
            )
            return response.choices[0].message.content
        except Exception as e:
            raise AIEndpointRequestError(f"Hugging Face request failed: {e}")