import base64
import os
from typing import List, Union, Literal, Type, Optional
from pydantic import BaseModel
import litellm
from litellm import acompletion
from ombra import step
from ombra.logger import logger

from .types import Model, Role, Message, ContentPart, LLMResponse, Usage

# Configure LiteLLM globally (e.g., enable caching, set default retries)
# For simplicity, we'll assume API keys are set as environment variables
# e.g., os.environ["GEMINI_API_KEY"], os.environ["OPENAI_API_KEY"]

@step()
async def call_llm(
    model: Union[Model, str],
    messages: List[Message],
    temperature: float = 0.7,
    max_tokens: int = 1024,
    retries: int = 3,
    timeout: int = 60,
    response_format: Optional[Union[dict, Type[BaseModel]]] = None,
    **kwargs
) -> LLMResponse:
    """
    Calls an LLM via LiteLLM, handling PDF conversion, retries, usage tracking, and structured outputs.
    This function is automatically checkpointed by Ombra.

    Args:
        model: The LLM model to use (e.g., Model.GEMINI_1_5_PRO or "gpt-4o").
        messages: A list of Message objects for the conversation.
        temperature: Sampling temperature.
        max_tokens: Maximum number of tokens to generate.
        retries: Number of retries for failed API calls.
        timeout: Timeout for the API call in seconds.
        response_format: Optional. Can be a dict (e.g. {"type": "json_object"}) or a Pydantic model class for structured outputs.
        kwargs: Additional keyword arguments to pass directly to litellm.acompletion.

    Returns:
        An LLMResponse object. If response_format was a Pydantic model, .content will be a JSON string matching that schema.
    """

    logger.info(f"🔍 call_llm received response_format: {response_format} (Type: {type(response_format)})")

    # Handle Pydantic model for structured outputs
    # We convert it to a schema dict that LiteLLM/Gemini can understand
    if response_format and isinstance(response_format, type) and issubclass(response_format, BaseModel):
        response_format = {
            "type": "json_object",
            "response_schema": response_format.model_json_schema()
        }

    # Prepare messages for LiteLLM
    litellm_messages = []
    for msg in messages:
        content_parts = []
        if isinstance(msg.content, str):
            content_parts.append({"type": "text", "text": msg.content})
        else:
            for part in msg.content:
                if isinstance(part, str):
                    content_parts.append({"type": "text", "text": part})
                elif isinstance(part, ContentPart):
                    if part.type == "text" and part.text:
                        content_parts.append({"type": "text", "text": part.text})
                    elif part.type == "file" and part.file:
                        # Handle PDF/File conversion to base64 for Gemini
                        file_bytes = await part.file.read()
                        base64_encoded_file = base64.b64encode(file_bytes).decode("utf-8")
                        mime_type = part.file.content_type if part.file.content_type else "application/octet-stream"
                        
                        # LiteLLM expects 'image_url' for base64 data URIs for multimodal input
                        # It handles mapping 'data:application/pdf;base64,' to Gemini's inlineData
                        content_parts.append({
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{base64_encoded_file}"
                            }
                        })
        
        # If there's only one text part, litellm can often take it directly as a string
        # Otherwise, send as a list of content parts
        final_content = content_parts[0]["text"] if len(content_parts) == 1 and content_parts[0]["type"] == "text" else content_parts

        litellm_messages.append({"role": msg.role.value, "content": final_content})

    try:
        response = await acompletion(
            model=model.value if isinstance(model, Model) else model,
            messages=litellm_messages,
            temperature=temperature,
            max_tokens=max_tokens,
            num_retries=retries, # LiteLLM parameter for retries
            timeout=timeout,     # LiteLLM parameter for timeout
            response_format=response_format,
            **kwargs
        )

        # Extract usage and content
        completion_content = response.choices[0].message.content if response.choices else ""
        usage_data = response.usage

        return LLMResponse(
            content=completion_content,
            role=Role.ASSISTANT,
            usage=Usage(
                prompt_tokens=usage_data.prompt_tokens,
                completion_tokens=usage_data.completion_tokens,
                total_tokens=usage_data.total_tokens,
            ),
            model=response.model,
            finish_reason=response.choices[0].finish_reason if response.choices else None
        )

    except litellm.exceptions.BadRequestError as e:
        # Handle specific LiteLLM errors, e.g., context window exceeded
        logger.error(f"LLM Bad Request Error: {e.message}")
        raise e
    except litellm.exceptions.RateLimitError as e:
        logger.warning(f"LLM Rate Limit Error: {e.message}")
        raise e
    except Exception as e:
        logger.error(f"An unexpected LLM error occurred: {e}")
        raise e
