from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union, cast

import litellm
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ModelResponse

from .transformation import GoogleGenAIAdapter

# Initialize adapter
GOOGLE_GENAI_ADAPTER = GoogleGenAIAdapter()


class GenerateContentToCompletionHandler:
    """Handler for transforming generate_content calls to completion format when provider config is None"""
    
    @staticmethod
    def _prepare_completion_kwargs(
        model: str,
        contents: Union[List[Dict[str, Any]], Dict[str, Any]],
        config: Optional[Dict[str, Any]] = None,
        stream: bool = False,
        litellm_params: Optional[GenericLiteLLMParams] = None,
        extra_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Prepare kwargs for litellm.completion/acompletion"""
        
        # Transform generate_content request to completion format
        completion_request = GOOGLE_GENAI_ADAPTER.translate_generate_content_to_completion(
            model=model,
            contents=contents,
            config=config,
            litellm_params=litellm_params,
            **(extra_kwargs or {})
        )
        
        completion_kwargs: Dict[str, Any] = dict(completion_request)
        
        if stream:
            completion_kwargs["stream"] = stream
                            
        return completion_kwargs

    @staticmethod
    async def async_generate_content_handler(
        model: str,
        contents: Union[List[Dict[str, Any]], Dict[str, Any]],
        litellm_params: GenericLiteLLMParams,
        config: Optional[Dict[str, Any]] = None,
        stream: bool = False,
        **kwargs,
    ) -> Union[Dict[str, Any], AsyncIterator[bytes]]:
        """Handle generate_content call asynchronously using completion adapter"""
        
        completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs(
            model=model,
            contents=contents,
            config=config,
            stream=stream,
            litellm_params=litellm_params,
            extra_kwargs=kwargs,
        )
        
        try:
            completion_response = await litellm.acompletion(**completion_kwargs)
            
            if stream:
                # Transform streaming completion response to generate_content format
                transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
                    completion_response
                )
                if transformed_stream is not None:
                    return transformed_stream
                raise ValueError("Failed to transform streaming response")
            else:
                # Transform completion response back to generate_content format
                generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
                    cast(ModelResponse, completion_response)
                )
                return generate_content_response
                
        except Exception as e:
            raise ValueError(
                f"Error calling litellm.acompletion for generate_content: {str(e)}"
            )

    @staticmethod
    def generate_content_handler(
        model: str,
        contents: Union[List[Dict[str, Any]], Dict[str, Any]],
        litellm_params: GenericLiteLLMParams,
        config: Optional[Dict[str, Any]] = None,
        stream: bool = False,
        _is_async: bool = False,
        **kwargs,
    ) -> Union[Dict[str, Any], AsyncIterator[bytes], Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]]]:
        """Handle generate_content call using completion adapter"""
        
        if _is_async:
            return GenerateContentToCompletionHandler.async_generate_content_handler(
                model=model,
                contents=contents,
                config=config,
                stream=stream,
                litellm_params=litellm_params,
                **kwargs,
            )
            
        completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs(
            model=model,
            contents=contents,
            config=config,
            stream=stream,
            litellm_params=litellm_params,
            extra_kwargs=kwargs,
        )
        
        try:
            completion_response = litellm.completion(**completion_kwargs)
            
            if stream:
                # Transform streaming completion response to generate_content format
                transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
                    completion_response
                )
                if transformed_stream is not None:
                    return transformed_stream
                raise ValueError("Failed to transform streaming response")
            else:
                # Transform completion response back to generate_content format
                generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
                    cast(ModelResponse, completion_response)
                )
                return generate_content_response
                
        except Exception as e:
            raise ValueError(
                f"Error calling litellm.completion for generate_content: {str(e)}"
            ) 