"""OpenAI Client Implementation"""
import asyncio
import random
from typing import AsyncIterator, Optional, Union, List, Dict, Any
from pathlib import Path
from openai import AsyncOpenAI
from openai._exceptions import APIError as OpenAIAPIError, APIConnectionError

from .base import BaseClient
from ..exceptions import APIError, NetworkError, ValidationError, RateLimitError
from ..utils import validate_prompt, create_openai_image_message
from ..models import GenerateResponse, StreamingResponse, TokenUsage


class OpenAIClient(BaseClient):
    """OpenAI async client"""
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
        max_retries: int = 2,
        rate_limit_retry_delay: float = 2.0,
        enable_rate_limit_backoff: bool = True
    ):
        """
        Initialize OpenAI client
        
        :param api_key: API key, default read from environment variable
        :param base_url: API base URL, default read from environment variable
        :param model: Model name, default read from environment variable
        :param max_retries: Maximum number of retries for rate limit errors, default 2
        :param rate_limit_retry_delay: Initial delay in seconds for rate limit retries, default 2.0
        :param enable_rate_limit_backoff: Enable exponential backoff for rate limit errors, default True
        """
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.max_retries = max_retries
        self.rate_limit_retry_delay = rate_limit_retry_delay
        self.enable_rate_limit_backoff = enable_rate_limit_backoff
        # Reduce default retries to avoid request accumulation during rate limiting
        self.client = AsyncOpenAI(api_key=api_key, base_url=base_url, max_retries=0)
    
    def add_image_to_messages(
        self,
        messages: List[Dict[str, Any]],
        image_path: Union[str, Path],
        text: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Add image to messages list for vision API
        
        :param messages: Existing message list
        :param image_path: Path to image file
        :param text: Optional text prompt to accompany the image
        :return: Updated messages list with image message
        """
        image_message = create_openai_image_message(image_path, text)
        messages.append(image_message)
        return messages
    
    async def generate(
        self,
        messages: List[Dict[str, Any]],
        stream: bool = False,
        temperature=0.6,
        **kwargs
    ) -> Union[GenerateResponse, StreamingResponse]:
        """
        Generate response with automatic rate limit handling and exponential backoff
        
        :param messages: Message list (can include image messages created with add_image_to_messages)
        :param stream: Whether to stream response
        :param temperature: Temperature parameter
        :return: StreamingResponse for streaming, GenerateResponse for non-streaming
        """
        # Build request parameters
        request_params = {
            'model': self.model,
            'messages': messages,
            'stream': stream,
            "temperature": temperature,
            **kwargs
        }
        
        # Retry logic with exponential backoff for rate limit errors
        last_exception = None
        for attempt in range(self.max_retries + 1):
            try:
                if stream:
                    response = await self.client.chat.completions.create(**request_params)
                    return await self._handle_streaming_response(response)
                else:
                    response = await self.client.chat.completions.create(**request_params)
                    content = response.choices[0].message.content or ""
                    print(response)

                    if response.usage:
                        usage = TokenUsage(
                            prompt_tokens=response.usage.prompt_tokens,
                            completion_tokens=response.usage.completion_tokens,
                            total_tokens=response.usage.total_tokens
                        )
                    
                    return GenerateResponse(content=content, usage=usage)
            
            except APIConnectionError as e:
                raise NetworkError(f"Network error when calling OpenAI API: {str(e)}") from e
            
            except OpenAIAPIError as e:
                # Check if it's a rate limit error (429)
                status_code = getattr(e, 'status_code', None)
                if status_code == 429:
                    last_exception = e
                    
                    # If this is the last attempt, raise RateLimitError
                    if attempt >= self.max_retries:
                        retry_after = self._extract_retry_after(e)
                        error_msg = (
                            f"Rate limit exceeded (429). "
                            f"Retried {self.max_retries} times. "
                            f"Last error: {str(e)}"
                        )
                        if retry_after:
                            error_msg += f" Retry-After: {retry_after} seconds"
                        raise RateLimitError(error_msg) from e
                    
                    # Calculate wait time with exponential backoff
                    wait_time = self._calculate_wait_time(e, attempt)
                    print(f"Rate limit error (429) encountered. Waiting {wait_time:.2f} seconds before retry {attempt + 1}/{self.max_retries}...")
                    
                    await asyncio.sleep(wait_time)
                    continue
                else:
                    # For non-429 errors, raise immediately
                    raise APIError(f"OpenAI API error: {str(e)}") from e
            
            except Exception as e:
                raise APIError(f"Unexpected error: {str(e)}") from e
        
        # This should never be reached, but just in case
        if last_exception:
            raise RateLimitError(f"Rate limit exceeded after {self.max_retries} retries: {str(last_exception)}") from last_exception
        raise APIError("Unexpected error: max retries exceeded without exception")
    
    def _extract_retry_after(self, error: OpenAIAPIError) -> Optional[float]:
        """
        Extract Retry-After value from error response headers
        
        :param error: OpenAI API error object
        :return: Retry-After value in seconds, or None if not found
        """
        try:
            # Check if error has response attribute with headers
            if hasattr(error, 'response') and error.response:
                headers = getattr(error.response, 'headers', None)
                if headers:
                    retry_after = headers.get('Retry-After') or headers.get('retry-after')
                    if retry_after:
                        try:
                            return float(retry_after)
                        except (ValueError, TypeError):
                            pass
        except Exception:
            pass
        return None
    
    def _calculate_wait_time(self, error: OpenAIAPIError, attempt: int) -> float:
        """
        Calculate wait time for rate limit retry with exponential backoff
        
        :param error: OpenAI API error object
        :param attempt: Current attempt number (0-indexed)
        :return: Wait time in seconds
        """
        # First, try to use Retry-After header if available
        retry_after = self._extract_retry_after(error)
        if retry_after:
            return retry_after
        
        # Otherwise, use exponential backoff
        if self.enable_rate_limit_backoff:
            # Exponential backoff: base_delay * (2 ^ attempt)
            # Add jitter to avoid thundering herd problem
            base_delay = self.rate_limit_retry_delay
            exponential_delay = base_delay * (2 ** attempt)
            # Add random jitter (0.5 to 1.5 times the delay)
            jitter = random.uniform(0.5, 1.5)
            wait_time = exponential_delay * jitter
            # Cap maximum wait time at 60 seconds
            return min(wait_time, 60.0)
        else:
            # Linear backoff
            return self.rate_limit_retry_delay * (attempt + 1)
    
    async def _handle_streaming_response(self, response) -> StreamingResponse:
        """
        Handle OpenAI streaming response
        
        :param response: OpenAI streaming response object
        :return: StreamingResponse object
        """
        usage = None
        
        async def stream_generator() -> AsyncIterator[str]:
            nonlocal usage
            async for chunk in response:
                if chunk.choices:
                    delta = chunk.choices[0].delta
                    if delta and delta.content:
                        yield delta.content

                if hasattr(chunk, 'usage'):
                    usage = TokenUsage(
                        prompt_tokens=chunk.usage["prompt_tokens"],
                        completion_tokens=chunk.usage["completion_tokens"],
                        total_tokens=chunk.usage["total_tokens"]
                    )


        async def get_usage():
            return usage
        
        return StreamingResponse(stream_generator(), usage_getter=get_usage)
    
    async def generate_with_tools(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto",
        temperature: float = 0.6,
        **kwargs
    ):
        """
        Generate response with tools support, returns raw OpenAI response object
        This method includes automatic rate limit handling and exponential backoff
        
        :param messages: Message list
        :param tools: List of tool schemas (optional)
        :param tool_choice: Tool choice parameter (default: "auto")
        :param temperature: Temperature parameter
        :param kwargs: Additional parameters to pass to API
        :return: Raw OpenAI response object (for accessing tool_calls, etc.)
        """
        # Build request parameters
        request_params = {
            'model': self.model,
            'messages': messages,
            'temperature': temperature,
            **kwargs
        }
        
        # Add tools if provided
        if tools is not None:
            request_params['tools'] = tools
        if tool_choice is not None:
            request_params['tool_choice'] = tool_choice
        
        # Retry logic with exponential backoff for rate limit errors
        last_exception = None
        for attempt in range(self.max_retries + 1):
            try:
                response = await self.client.chat.completions.create(**request_params)
                return response
            
            except APIConnectionError as e:
                raise NetworkError(f"Network error when calling OpenAI API: {str(e)}") from e
            
            except OpenAIAPIError as e:
                # Check if it's a rate limit error (429)
                status_code = getattr(e, 'status_code', None)
                if status_code == 429:
                    last_exception = e
                    
                    # If this is the last attempt, raise RateLimitError
                    if attempt >= self.max_retries:
                        retry_after = self._extract_retry_after(e)
                        error_msg = (
                            f"Rate limit exceeded (429). "
                            f"Retried {self.max_retries} times. "
                            f"Last error: {str(e)}"
                        )
                        if retry_after:
                            error_msg += f" Retry-After: {retry_after} seconds"
                        raise RateLimitError(error_msg) from e
                    
                    # Calculate wait time with exponential backoff
                    wait_time = self._calculate_wait_time(e, attempt)
                    print(f"Rate limit error (429) encountered. Waiting {wait_time:.2f} seconds before retry {attempt + 1}/{self.max_retries}...")
                    
                    await asyncio.sleep(wait_time)
                    continue
                else:
                    # For non-429 errors, raise immediately
                    raise APIError(f"OpenAI API error: {str(e)}") from e
            
            except Exception as e:
                raise APIError(f"Unexpected error: {str(e)}") from e
        
        # This should never be reached, but just in case
        if last_exception:
            raise RateLimitError(f"Rate limit exceeded after {self.max_retries} retries: {str(last_exception)}") from last_exception
        raise APIError("Unexpected error: max retries exceeded without exception")


