"""
OpenAI Provider Module

Implements LLM interface for OpenAI API:
- GPT-4, GPT-4o, GPT-3.5-turbo
- Tool/function calling
- Streaming support
"""

import json
import time
from typing import Optional

from openai import OpenAI, APIError, RateLimitError as OpenAIRateLimitError, AuthenticationError as OpenAIAuthError

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


@register_provider("openai")
class OpenAIProvider(LLMProvider):
    """
    OpenAI API provider implementation.
    
    Supports:
    - GPT-4o, GPT-4-turbo, GPT-4, GPT-3.5-turbo
    - Function/tool calling
    - JSON mode
    """

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        self.client = OpenAI(
            api_key=config.api_key,
            base_url=config.api_base,
            timeout=config.timeout,
        )

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

        kwargs = {
            "model": self.config.model,
            "messages": formatted_messages,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
        }

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

        for attempt in range(self.config.max_retries):
            try:
                response = self.client.chat.completions.create(**kwargs)
                return self._parse_response(response)

            except OpenAIRateLimitError as e:
                if attempt < self.config.max_retries - 1:
                    wait_time = 2 ** attempt
                    time.sleep(wait_time)
                    continue
                raise RateLimitError(str(e))

            except OpenAIAuthError as e:
                raise AuthenticationError(f"OpenAI authentication failed: {e}")

            except APIError as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError(f"OpenAI API error: {e}")

        raise LLMError("Max retries exceeded")

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

        for msg in messages:
            if msg.role == "tool":
                # Tool response message
                formatted.append({
                    "role": "tool",
                    "content": msg.content,
                    "tool_call_id": msg.tool_call_id or "",
                })
            elif msg.tool_calls:
                # Assistant message with tool calls
                formatted.append({
                    "role": "assistant",
                    "content": msg.content or "",
                    "tool_calls": msg.tool_calls,
                })
            else:
                # Regular message
                formatted.append({
                    "role": msg.role,
                    "content": msg.content,
                })

        return formatted

    def _parse_response(self, response) -> CompletionResponse:
        """Parse OpenAI API response."""
        choice = response.choices[0]
        message = choice.message

        tool_calls = []
        if message.tool_calls:
            for tc in message.tool_calls:
                try:
                    args = json.loads(tc.function.arguments)
                except json.JSONDecodeError:
                    args = {}

                tool_calls.append(ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=args,
                ))

        usage = None
        if response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }

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

    def count_tokens(self, text: str) -> int:
        """
        Estimate token count using tiktoken if available.
        Falls back to character-based estimation.
        """
        try:
            import tiktoken
            encoding = tiktoken.encoding_for_model(self.config.model)
            return len(encoding.encode(text))
        except ImportError:
            # Fallback: ~4 chars per token
            return len(text) // 4
