"""
HuggingFace Provider Module

HuggingFace Inference API for accessing models hosted on HuggingFace:
- Meta Llama models
- Mistral models
- And many other open-source models
"""

import json
import time
from typing import Optional

import httpx

from ai_coder.llm.interface import (
    LLMProvider,
    LLMConfig,
    Message,
    ToolCall,
    CompletionResponse,
    LLMError,
    RateLimitError,
    AuthenticationError,
    register_provider,
)


@register_provider("huggingface")
class HuggingFaceProvider(LLMProvider):
    """
    HuggingFace Inference API provider.
    
    Get your API key at: https://huggingface.co/settings/tokens
    
    Popular models:
    - meta-llama/Llama-3.3-70B-Instruct
    - meta-llama/Llama-3.1-8B-Instruct
    - mistralai/Mistral-7B-Instruct-v0.3
    - microsoft/Phi-3-mini-4k-instruct
    """

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        
        self.api_key = config.api_key
        if not self.api_key:
            raise AuthenticationError("HUGGINGFACE_API_KEY is required. Get one at https://huggingface.co/settings/tokens")
        
        # HuggingFace Router API endpoint (model is specified in request body, not URL)
        self.base_url = config.api_base or "https://router.huggingface.co/v1"
        
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        
        self.client = httpx.Client(timeout=config.timeout)

    def complete(
        self,
        messages: list[Message],
        tools: Optional[list[dict]] = None,
    ) -> CompletionResponse:
        """Generate completion using HuggingFace Inference API."""
        formatted_messages = self._format_messages(messages)

        payload = {
            "model": self.config.model,
            "messages": formatted_messages,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
            "stream": False,
        }

        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"

        url = f"{self.base_url}/chat/completions"

        for attempt in range(self.config.max_retries):
            try:
                response = self.client.post(url, headers=self.headers, json=payload)
                
                if response.status_code == 401:
                    raise AuthenticationError("HuggingFace API key is invalid")
                
                if response.status_code == 429:
                    if attempt < self.config.max_retries - 1:
                        wait_time = 2 ** attempt
                        time.sleep(wait_time)
                        continue
                    raise RateLimitError("Rate limit exceeded")
                
                if response.status_code == 503:
                    # Model is loading - wait and retry
                    if attempt < self.config.max_retries - 1:
                        wait_time = 20  # HuggingFace models can take time to load
                        time.sleep(wait_time)
                        continue
                    raise LLMError("Model is still loading. Please try again.")
                
                if response.status_code >= 400:
                    error_text = response.text[:500]
                    raise LLMError(f"HuggingFace API error ({response.status_code}): {error_text}")
                
                return self._parse_response(response.json())

            except httpx.TimeoutException:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError("Request timed out")
            
            except httpx.RequestError as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError(f"Connection error: {e}")

        raise LLMError("Max retries exceeded")

    def _format_messages(self, messages: list[Message]) -> list[dict]:
        """Format messages for HuggingFace API (OpenAI-compatible)."""
        formatted = []

        for msg in messages:
            if msg.role == "tool":
                formatted.append({
                    "role": "tool",
                    "content": msg.content,
                    "tool_call_id": msg.tool_call_id or "",
                })
            elif msg.tool_calls:
                formatted.append({
                    "role": "assistant",
                    "content": msg.content or "",
                    "tool_calls": msg.tool_calls,
                })
            else:
                formatted.append({
                    "role": msg.role,
                    "content": msg.content,
                })

        return formatted

    def _parse_response(self, data: dict) -> CompletionResponse:
        """Parse HuggingFace API response."""
        choices = data.get("choices", [])
        if not choices:
            return CompletionResponse(content="No response generated")

        choice = choices[0]
        message = choice.get("message", {})

        tool_calls = []
        raw_tool_calls = message.get("tool_calls", [])
        for tc in raw_tool_calls:
            try:
                func = tc.get("function", {})
                args = func.get("arguments", "{}")
                if isinstance(args, str):
                    args = json.loads(args)
                
                tool_calls.append(ToolCall(
                    id=tc.get("id", ""),
                    name=func.get("name", ""),
                    arguments=args,
                ))
            except json.JSONDecodeError:
                pass

        usage = None
        if "usage" in data:
            usage = {
                "prompt_tokens": data["usage"].get("prompt_tokens", 0),
                "completion_tokens": data["usage"].get("completion_tokens", 0),
                "total_tokens": data["usage"].get("total_tokens", 0),
            }

        return CompletionResponse(
            content=message.get("content"),
            tool_calls=tool_calls,
            finish_reason=choice.get("finish_reason", "stop"),
            usage=usage,
            raw_response=data,
        )

    def count_tokens(self, text: str) -> int:
        """Estimate token count."""
        return len(text) // 4
