import base64
import json
import os
import tempfile
from typing import List, Optional, Sequence, Type, Union

from arbol import aprint
from openai import OpenAI
from pydantic import BaseModel

from litemind.agent.messages.message import Message
from litemind.agent.messages.message_block import MessageBlock
from litemind.agent.tools.base_tool import BaseTool
from litemind.agent.tools.builtin_tools.mcp_tool import BuiltinMCPTool
from litemind.agent.tools.builtin_tools.web_search_tool import BuiltinWebSearchTool
from litemind.agent.tools.toolset import ToolSet
from litemind.apis.base_api import ModelFeatures
from litemind.apis.callbacks.api_callback_manager import ApiCallbackManager
from litemind.apis.default_api import DefaultApi
from litemind.apis.exceptions import APIError, APINotAvailableError
from litemind.apis.model_registry import get_default_model_registry
from litemind.apis.providers.openai.utils.check_availability import (
    check_openai_api_availability,
)
from litemind.apis.providers.openai.utils.is_reasoning import (
    does_openai_model_support_temperature,
    is_openai_reasoning_model,
)
from litemind.apis.providers.openai.utils.list_models import (
    _get_raw_openai_model_list,
    get_openai_model_list,
)
from litemind.media.media_base import MediaBase
from litemind.media.types.media_action import Action
from litemind.media.types.media_image import Image
from litemind.media.types.media_text import Text
from litemind.utils.normalise_uri_to_local_file_path import uri_to_local_file_path


class OpenAIApi(DefaultApi):
    """
    An OpenAI Response API implementation conforming to the BaseApi interface.

    This implementation uses OpenAI's newer Response API which is stateful and provides
    built-in tools like web search, file search, and computer use. It's designed to be
    simpler and more flexible than the chat completions API.

    Key differences from chat completions:
    - Uses client.responses.create() instead of client.chat.completions.create()
    - Stateful - maintains conversation history automatically
    - Built-in tools like web_search, file_search
    - Different request/response format with 'input' parameter
    - Can chain responses using previous_response_id

    Set the OPENAI_API_KEY environment variable to your OpenAI API key, or pass it as an argument to the constructor.
    You can get your API key from https://platform.openai.com/account/api-keys.

    Note: This implementation requires models that support the Response API.
    Not compatible with older models that only support chat completions.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        allow_media_conversions: bool = True,
        allow_media_conversions_with_models: bool = True,
        callback_manager: Optional[ApiCallbackManager] = None,
        **kwargs,
    ):
        """
        Initialize the OpenAI Response API client.

        Parameters
        ----------
        api_key: Optional[str]
            The API key for OpenAI. If not provided, we'll read from OPENAI_API_KEY env var.
        base_url: Optional[str]
            The base URL for the OpenAI or compatible API.
        allow_media_conversions: bool
            If True, the API will allow media conversions using the default media converter.
        allow_media_conversions_with_models: bool
            If True, the API will allow media conversions using models that support the required features.
        kwargs: dict
            Additional options passed to `OpenAI(...)`.
        """

        super().__init__(
            allow_media_conversions=allow_media_conversions,
            allow_media_conversions_with_models=allow_media_conversions_with_models,
            callback_manager=callback_manager,
        )

        # get key from environmental variables:
        if api_key is None:
            api_key = os.environ.get("OPENAI_API_KEY")

        # raise error if api_key is None:
        if api_key is None:
            raise APIError(
                "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable"
            )

        # Save the API key:
        self.api_key = api_key

        # Save base URL:
        self.base_url = base_url

        # Save additional kwargs:
        self.kwargs = kwargs

        try:
            # Initialize the model registry for feature lookups:
            self.model_registry = get_default_model_registry()

            # ensure the timeout parameter is set in kwargs and high enough:
            if "timeout" not in self.kwargs or self.kwargs["timeout"] < 6000:
                self.kwargs["timeout"] = 6000

            # Create an OpenAI client:
            self.client: OpenAI = OpenAI(
                api_key=self.api_key, base_url=self.base_url, **self.kwargs
            )

            # Get the raw list of models:
            self._raw_model_list = _get_raw_openai_model_list(self.client)

            # Get the list of models from OpenAI (filter to only Response API compatible models):
            self._model_list = get_openai_model_list(self._raw_model_list)

        except Exception as e:
            # Print stack trace:
            import traceback

            traceback.print_exc()
            raise APINotAvailableError(
                f"Error initializing OpenAI Response API client: {e}"
            )

    def check_availability_and_credentials(self, api_key: Optional[str] = None) -> bool:
        """Check if the API is available and credentials are valid."""

        # Check if the API key is provided:
        if api_key is not None:
            client = OpenAI(api_key=api_key, base_url=self.base_url, **self.kwargs)
        else:
            client = self.client

        # Check the OpenAI API key:
        result = check_openai_api_availability(client)

        # Call the callback manager:
        self.callback_manager.on_availability_check(result)

        return result

    def list_models(
        self,
        features: Optional[
            Union[str, List[str], ModelFeatures, Sequence[ModelFeatures]]
        ] = None,
        non_features: Optional[
            Union[str, List[str], ModelFeatures, Sequence[ModelFeatures]]
        ] = None,
        media_types: Optional[Sequence[Type[MediaBase]]] = None,
    ) -> List[str]:
        """List available models that support the Response API."""

        try:
            # Normalise the features:
            features = ModelFeatures.normalise(features)
            non_features = ModelFeatures.normalise(non_features)

            # Get the Response API compatible models:
            model_list = list(self._model_list)

            # Add the models from the super class:
            model_list += super().list_models()

            # Filter the models based on the features:
            if features:
                model_list = self._filter_models(
                    model_list,
                    features=features,
                    non_features=non_features,
                    media_types=media_types,
                )

            # Call callbacks:
            self.callback_manager.on_model_list(model_list)

            return model_list

        except Exception:
            raise APIError("Error fetching model list from OpenAI Response API.")

    def get_best_model(
        self,
        features: Optional[
            Union[str, List[str], ModelFeatures, Sequence[ModelFeatures]]
        ] = None,
        non_features: Optional[
            Union[str, List[str], ModelFeatures, Sequence[ModelFeatures]]
        ] = None,
        media_types: Optional[Sequence[Type[MediaBase]]] = None,
        exclusion_filters: Optional[Union[str, List[str]]] = None,
    ) -> Optional[str]:
        """Get the best model that supports the required features."""

        # Normalise the features:
        features = ModelFeatures.normalise(features)
        non_features = ModelFeatures.normalise(non_features)

        # Get model list:
        model_list = self.list_models()

        # Filter the models based on the requirements:
        model_list = self._filter_models(
            model_list,
            features=features,
            non_features=non_features,
            media_types=media_types,
            exclusion_filters=exclusion_filters,
        )

        # By default, result is None:
        best_model = None

        if len(model_list) > 0:
            best_model = model_list[0]

        # Call the callback manager:
        self.callback_manager.on_best_model_selected(best_model)

        return best_model

    def has_model_support_for(
        self,
        features: Union[str, List[str], ModelFeatures, Sequence[ModelFeatures]],
        media_types: Optional[Sequence[Type[MediaBase]]] = None,
        model_name: Optional[str] = None,
    ) -> bool:
        """Check if a model supports the given features."""

        # Normalise the features:
        features = ModelFeatures.normalise(features)

        # Get the best model if not provided:
        if model_name is None:
            model_name = self.get_best_model(features=features, media_types=media_types)

        # If model_name is None then we return False:
        if model_name is None:
            return False

        # We check if the superclass says that the model supports the features:
        if super().has_model_support_for(
            features=features, media_types=media_types, model_name=model_name
        ):
            return True

        # Check if model is Response API compatible
        if model_name not in self._model_list:
            return False

        for feature in features:
            if not self.model_registry.supports_feature(
                self.__class__, model_name, feature
            ):
                return False

        return True

    def max_num_input_tokens(self, model_name: Optional[str] = None) -> int:
        """Get the maximum number of input tokens for the given model."""
        if model_name is None:
            model_name = self.get_best_model()

        # Try registry first (most accurate source)
        model_info = self.model_registry.get_model_info(self.__class__, model_name)
        if model_info and model_info.context_window:
            return model_info.context_window

        # Fallback for unknown models
        return 128_000

    def max_num_output_tokens(self, model_name: Optional[str] = None) -> int:
        """Get the maximum number of output tokens for the given model."""
        if model_name is None:
            model_name = self.get_best_model()

        # Try registry first (most accurate source)
        model_info = self.model_registry.get_model_info(self.__class__, model_name)
        if model_info and model_info.max_output_tokens:
            return model_info.max_output_tokens

        # Fallback for unknown models
        return 16_384

    def _convert_messages_to_response_input(
        self, messages: List[Message], is_tool_followup: bool = False
    ) -> Union[str, List[dict]]:
        """
        Convert Litemind messages to Response API input format.

        The Response API can accept either:
        1. A simple string for single-turn interactions
        2. An array of conversation items for multi-turn conversations

        Parameters
        ----------
        messages : List[Message]
            The messages to convert
        is_tool_followup : bool
            True if this is a follow-up call with tool results in the same Response API session.
            When False, tool messages are converted to conversation items.
            When True, this method should not be called (use function_call_outputs directly).
        """

        if len(messages) == 1 and len(messages[0].blocks) == 1:
            # Simple case: single message with text only
            if messages[0].blocks[0].has_type(Text):
                return messages[0].blocks[0].get_content()

        # Complex case: convert to conversation items format
        conversation_items = []

        for message in messages:
            if message.role == "system":
                # System messages are converted to user messages with special formatting
                content = self._convert_message_blocks_to_response_content(
                    message.blocks, is_input=True, is_system=True
                )

                if content:  # Only add if there's content
                    conversation_items.append(
                        {"type": "message", "role": "user", "content": content}
                    )

            elif message.role == "user":
                content = self._convert_message_blocks_to_response_content(
                    message.blocks, is_input=True
                )

                if content:  # Only add if there's content
                    conversation_items.append(
                        {"type": "message", "role": "user", "content": content}
                    )

            elif message.role == "assistant":
                # Assistant messages from previous turns
                content = self._convert_message_blocks_to_response_content(
                    message.blocks, is_input=False
                )

                if content:  # Only add if there's content
                    conversation_items.append(
                        {"type": "message", "role": "assistant", "content": content}
                    )

            elif message.role == "tool":
                # Tool use results - convert to user messages describing what happened
                # DO NOT convert to function_call_outputs when starting a new conversation

                if not is_tool_followup:
                    # Convert tool results to descriptive user messages for conversation history
                    content = []
                    for block in message.blocks:
                        if block.has_type(Action):
                            action = block.get_content()
                            if hasattr(action, "tool_name") and hasattr(
                                action, "result"
                            ):
                                tool_result_text = f"[Tool {action.tool_name} returned: {action.result}]"
                                content.append(
                                    {"type": "input_text", "text": tool_result_text}
                                )

                    if content:  # Only add if there's content
                        conversation_items.append(
                            {
                                "type": "message",
                                "role": "user",  # Tool results become user messages in conversation history
                                "content": content,
                            }
                        )

        return conversation_items

    def _convert_message_blocks_to_response_content(
        self, blocks: List[MessageBlock], is_input: bool = True, is_system: bool = False
    ) -> List[dict]:
        """
        Convert message blocks to Response API content format.

        Parameters
        ----------
        blocks : List[MessageBlock]
            The message blocks to convert
        is_input : bool
            True for input content (user messages), False for output content (assistant messages)
        is_system : bool
            True if this is a system message (adds [SYSTEM] prefix to text)
        """
        from litemind.media.types.media_audio import Audio
        from litemind.media.types.media_video import Video

        content = []

        for block in blocks:
            media = block.media

            if block.has_type(Text):
                text_content = media.get_content()
                if is_system:
                    text_content = f"[SYSTEM]: {text_content}"

                content.append(
                    {
                        "type": "input_text" if is_input else "output_text",
                        "text": text_content,
                    }
                )

            elif block.has_type(Image):
                if is_input:
                    # Use existing to_remote_or_data_uri method
                    image_uri = media.to_remote_or_data_uri()
                    content.append({"type": "input_image", "image_url": image_uri})
                else:
                    # For assistant output, describe the image
                    filename = (
                        media.get_filename()
                        if hasattr(media, "get_filename")
                        else "image"
                    )
                    content.append(
                        {"type": "output_text", "text": f"[Image: {filename}]"}
                    )

            elif block.has_type(Audio):
                if is_input:
                    # OpenAI Responses API doesn't support input_audio content type.
                    # Transcribe audio first using Whisper, then send as text.
                    try:
                        audio_uri = media.uri
                        transcription = self.transcribe_audio(audio_uri)
                        content.append(
                            {
                                "type": "input_text",
                                "text": f"[Audio transcription]: {transcription}",
                            }
                        )
                    except Exception:
                        # Fallback to text placeholder if transcription fails
                        content.append(
                            {
                                "type": "input_text",
                                "text": f"[Audio file: {media.get_filename()}]",
                            }
                        )
                else:
                    # For assistant output, describe the audio
                    content.append(
                        {
                            "type": "output_text",
                            "text": f"[Audio: {media.get_filename()}]",
                        }
                    )

            elif block.has_type(Video):
                # Graceful fallback for Video blocks that couldn't be converted
                filename = (
                    media.get_filename() if hasattr(media, "get_filename") else "video"
                )
                if is_input:
                    content.append(
                        {"type": "input_text", "text": f"[Video file: {filename}]"}
                    )
                else:
                    content.append(
                        {"type": "output_text", "text": f"[Video: {filename}]"}
                    )

            elif block.has_type(Action):
                # Handle tool calls for assistant messages
                if not is_input:
                    action = media.get_content()
                    if hasattr(action, "tool_name") and hasattr(action, "arguments"):
                        # This is a tool call - represent it as text in the conversation
                        tool_call_text = f"[Called tool {action.tool_name} with arguments: {action.arguments}]"
                        content.append({"type": "output_text", "text": tool_call_text})
                # For input (user) messages, Action blocks shouldn't normally appear

            else:
                # throw an error if unsupported media type:
                raise ValueError(
                    f"Unsupported media type in message block: {type(media).__name__}"
                )

        return content

    def _convert_response_to_messages(
        self, response, response_format: Optional[BaseModel] = None
    ) -> List[Message]:
        """Convert Response API output to Litemind messages with proper structured output handling."""

        messages = []

        for output_item in response.output:
            if output_item.type == "message" and output_item.role == "assistant":
                message = Message(role="assistant")

                for content_item in output_item.content:
                    if content_item.type == "output_text":
                        text_content = content_item.text

                        # Handle structured output - if response_format is specified and we get JSON text
                        if response_format and text_content.strip().startswith("{"):
                            try:
                                json_data = json.loads(text_content)
                                obj_instance = response_format(**json_data)
                                message.append_object(obj_instance)
                            except (json.JSONDecodeError, ValueError, TypeError) as e:
                                # If parsing fails, treat as regular text but also log the issue
                                print(
                                    f"Warning: Failed to parse structured output: {e}"
                                )
                                message.append_text(text_content)
                        else:
                            # Regular text content
                            message.append_text(text_content)

                messages.append(message)

            elif output_item.type == "function_call":
                # Handle custom function calls
                message = Message(role="assistant")
                message.append_tool_call(
                    tool_name=output_item.name,
                    arguments=(
                        json.loads(output_item.arguments)
                        if isinstance(output_item.arguments, str)
                        else output_item.arguments
                    ),
                    id=output_item.call_id,
                )
                messages.append(message)

            # Handle built-in tool calls (web search, file search, etc.)
            elif output_item.type in [
                "web_search_call",
                "file_search_call",
                "computer_use_call",
            ]:
                # Built-in tools are handled automatically by the Response API
                # We don't need to create messages for these
                pass

        return messages

    def _format_tools_for_response_api(self, toolset: ToolSet) -> List[dict]:
        """Format tools for the Response API."""
        tools = []

        for tool in toolset.list_tools():

            # Skip built-in tools as they are handled automatically by the Response API
            if tool.is_builtin():
                continue

            # Response API uses a slightly different format than Chat Completions API
            tool_def = {
                "type": "function",
                "name": tool.name,
                "description": tool.description,
                "parameters": tool.arguments_schema,
            }
            tools.append(tool_def)

        return tools

    def generate_text(
        self,
        messages: List[Message],
        model_name: Optional[str] = None,
        temperature: Optional[float] = 0.0,
        max_num_output_tokens: Optional[int] = None,
        toolset: Optional[ToolSet] = None,
        use_tools: bool = True,
        response_format: Optional[BaseModel] = None,
        **kwargs,
    ) -> List[Message]:
        """Generate text using the Response API.

        Parameters
        ----------
        messages : List[Message]
            The input messages for the conversation.
        model_name : Optional[str]
            The name of the model to use.
        temperature : Optional[float]
            The temperature for text generation.
        max_num_output_tokens : Optional[int]
            Maximum number of output tokens.
        toolset : Optional[ToolSet]
            Custom tools to make available to the model.
        use_tools : bool
            Whether to use tools (both custom and built-in).
        response_format : Optional[BaseModel]
            Pydantic model for structured output.
        **kwargs
            Additional parameters to pass to the API.

        Returns
        -------
        List[Message]
            The generated response messages.
        """

        use_file_search: bool = False
        vector_store_ids: Optional[List[str]] = None
        use_computer_use: bool = False
        stream: bool = kwargs["stream"] if "stream" in kwargs else True

        # Validate inputs using parent method:
        super().generate_text(
            messages=messages,
            model_name=model_name,
            temperature=temperature,
            max_num_output_tokens=max_num_output_tokens,
            toolset=toolset,
            use_tools=use_tools,
            response_format=response_format,
            **kwargs,
        )

        # Set default model if not provided
        if model_name is None:
            model_name = self._get_best_model_for_text_generation(
                messages, toolset if use_tools else None, response_format
            )

        # **WORKAROUND** for deep research models that don't like streaming:
        if model_name and "deep-research" in model_name:
            stream = False
            aprint(
                "Deep Research models do not support streaming. Disabling streaming for this request."
            )

        # Handle model-specific parameter restrictions
        reasoning_effort = None
        is_reasoning_model = is_openai_reasoning_model(model_name)

        # Check if model supports temperature (o-series and gpt-5.x don't)
        if not does_openai_model_support_temperature(model_name):
            temperature = None

        if is_reasoning_model:
            # Extract reasoning effort from model name (e.g., "o3-high" -> "high")
            reasoning_effort_suffix = model_name.split("-")[-1]

            # O-series models support: low, medium, high
            if reasoning_effort_suffix in ["low", "medium", "high"]:
                reasoning_effort = reasoning_effort_suffix
                # Remove the suffix from the model name
                model_name_native = model_name.replace(
                    f"-{reasoning_effort_suffix}", ""
                )
            else:
                reasoning_effort = None
                model_name_native = model_name
        else:
            model_name_native = model_name

        # Preprocess messages:
        preprocessed_messages = self._preprocess_messages(
            messages=messages,
            allowed_media_types=self._get_allowed_media_types_for_text_generation(
                model_name=model_name
            ),
        )

        # Get max num of output tokens for model if not provided:
        if max_num_output_tokens is None:
            max_num_output_tokens = self.max_num_output_tokens(model_name)

        # Prepare tools (both custom and built-in)
        tools = []
        if use_tools:
            # Add custom function tools
            if toolset:
                tools.extend(self._format_tools_for_response_api(toolset))

                for builtinTool in toolset.list_builtin_tools():
                    if isinstance(builtinTool, BuiltinWebSearchTool):
                        web_search_tool: BuiltinWebSearchTool = builtinTool
                        # we can extract parameters from the web search tool:
                        search_context_size = web_search_tool.search_context_size
                        tools.append(
                            {
                                "type": "web_search_preview",
                                "search_context_size": search_context_size,
                            }
                        )
                    elif isinstance(builtinTool, BuiltinMCPTool):
                        mcp_tool: BuiltinMCPTool = builtinTool
                        # we can extract parameters from the MCP tool:
                        tools.append(
                            {
                                "type": "mcp",
                                "server_label": mcp_tool.server_name,
                                "server_url": mcp_tool.server_url,
                                "headers": mcp_tool.headers,
                                "require_approval": "never",
                                "allowed_tools": mcp_tool.allowed_tools,
                            }
                        )
                    else:
                        # Throw an exception if unknown built-in tool is used:
                        raise ValueError(
                            "Unknown built-in tool: {}".format(builtinTool.name)
                        )

            if use_file_search:
                if not vector_store_ids:
                    raise ValueError(
                        "vector_store_ids must be provided when use_file_search=True"
                    )
                tools.append(
                    {"type": "file_search", "vector_store_ids": vector_store_ids}
                )

            if use_computer_use:
                tools.append({"type": "computer_use"})

        try:
            # List of new messages to append to the response:
            new_messages = []
            previous_response_id = None  # Track for stateful API
            result_messages = []

            # Loop until we get a response that doesn't require tool use:
            while True:
                if previous_response_id is None:
                    # Initial request: convert full conversation history
                    response_input = self._convert_messages_to_response_input(
                        preprocessed_messages, is_tool_followup=False
                    )

                    # Prepare initial API parameters
                    api_params = {
                        "model": model_name_native,
                        "input": response_input,
                        "max_output_tokens": max_num_output_tokens,
                    }

                else:
                    # Subsequent request: use previous_response_id with tool outputs
                    # response_input should already be set to function_call_outputs from previous iteration
                    api_params = {
                        "model": model_name_native,
                        "input": response_input,  # This will be function_call_outputs
                        "previous_response_id": previous_response_id,
                        "max_output_tokens": max_num_output_tokens,
                    }

                # Add structured output support - THIS IS THE KEY FIX
                if response_format:
                    api_params["text_format"] = response_format

                # Only add temperature if it's not None (o-series models don't support it)
                if temperature is not None:
                    api_params["temperature"] = temperature

                # Add reasoning effort if specified (for o-series models)
                if reasoning_effort is not None:
                    api_params["reasoning"] = {
                        "effort": reasoning_effort,
                        "summary": "auto",
                    }

                # Only add tools if there are any
                if tools:
                    api_params["tools"] = tools

                # Add any additional kwargs
                api_params.update(kwargs)

                # Create response using Response API with streaming
                try:
                    if stream:
                        # Streaming request
                        with self.client.responses.stream(
                            **api_params
                        ) as streaming_response:
                            # Process streaming events - based on official OpenAI example
                            try:
                                for event in streaming_response:
                                    if "output_text" in event.type:
                                        # Extract the text delta for streaming callback
                                        if hasattr(event, "delta"):
                                            self.callback_manager.on_text_streaming(
                                                fragment=event.delta, **kwargs
                                            )
                                        elif hasattr(event, "text"):
                                            self.callback_manager.on_text_streaming(
                                                fragment=event.text, **kwargs
                                            )
                            except Exception as e:
                                # Handle any errors during streaming
                                aprint(f"Unexpected error during streaming: {e}")
                                # print stack trace:
                                import traceback

                                traceback.print_exc()

                            # Get the final response - identical to non-streaming!
                            response = streaming_response.get_final_response()
                    else:
                        # Non-streaming request
                        response = self.client.responses.create(**api_params)

                except AttributeError:
                    # Fallback if streaming method not available
                    response = self.client.responses.create(**api_params)

                # Convert response to Litemind messages
                result_messages = self._convert_response_to_messages(response)

                # Append response message to original, preprocessed, and new messages:
                if result_messages:
                    messages.extend(result_messages)
                    preprocessed_messages.extend(result_messages)
                    new_messages.extend(result_messages)

                # Check if we have tool calls to execute
                tool_calls = []
                for output_item in response.output:
                    if output_item.type == "function_call":
                        tool_calls.append(output_item)

                if tool_calls and use_tools:
                    # Create tool use message to hold results for our internal tracking
                    tool_use_message = Message(role="tool")
                    messages.append(tool_use_message)
                    preprocessed_messages.append(tool_use_message)
                    new_messages.append(tool_use_message)

                    # Prepare function call outputs for next Response API call
                    function_call_outputs = []

                    # Execute each tool call
                    for function_call in tool_calls:
                        tool_name = function_call.name
                        tool_arguments = (
                            json.loads(function_call.arguments)
                            if isinstance(function_call.arguments, str)
                            else function_call.arguments
                        )
                        call_id = function_call.call_id

                        # Get the corresponding tool in toolset:
                        tool: Optional[BaseTool] = (
                            toolset.get_tool(tool_name) if toolset else None
                        )

                        if tool:
                            try:
                                # Execute the tool
                                result = tool(**tool_arguments)

                                # If not a string, convert from JSON:
                                if not isinstance(result, str):
                                    result = json.dumps(result, default=str)

                            except Exception as e:
                                result = f"Function '{tool_name}' error: {e}"
                        else:
                            result = f"(Tool '{tool_name}' use requested, but tool not found.)"

                        # Append the tool call result to our internal message:
                        tool_use_message.append_tool_use(
                            tool_name=tool_name,
                            arguments=tool_arguments,
                            result=result,
                            id=call_id,
                        )

                        # Prepare function call output for Response API
                        function_call_outputs.append(
                            {
                                "type": "function_call_output",
                                "call_id": call_id,
                                "output": result,
                            }
                        )

                    # Set up for next iteration with tool results
                    response_input = function_call_outputs
                    previous_response_id = response.id
                    # Continue the loop to send tool results and potentially get more tool calls

                else:
                    # No tool calls, we're done
                    break

            # If using structured output, convert the text to object
            if response_format and result_messages:
                last_message = result_messages[-1]
                if last_message.blocks and last_message.blocks[-1].has_type(Text):
                    text_content = last_message.blocks[-1].get_content()
                    try:
                        obj_data = json.loads(text_content)
                        obj_instance = response_format(**obj_data)
                        # Replace the Text block with the Object block:
                        from litemind.media.types.media_object import Object

                        last_message.blocks[-1] = MessageBlock(
                            media=Object(obj_instance)
                        )
                    except (json.JSONDecodeError, ValueError):
                        # If parsing fails, keep as text
                        pass

            # Add parameters to kwargs for callback:
            kwargs.update(
                {
                    "model_name": model_name,
                    "temperature": temperature,
                    "max_output_tokens": max_num_output_tokens,
                    "toolset": toolset,
                    "response_format": response_format,
                    "stream": stream,
                }
            )

            # Call the callback manager:
            if result_messages:
                self.callback_manager.on_text_generation(
                    messages=messages, response=result_messages[-1], **kwargs
                )

            return new_messages

        except Exception as e:
            import traceback

            traceback.print_exc()
            raise APIError(f"OpenAI Response API generate text error: {e}")

    def transcribe_audio(
        self, audio_uri: str, model_name: Optional[str] = None, **kwargs
    ) -> str:
        """Transcribe audio using OpenAI's Whisper API."""

        # If no model name is provided, use the best model:
        if model_name is None:
            model_name = self.get_best_model(features=ModelFeatures.AudioTranscription)

        # If no model is available then delegate to superclass:
        if model_name is None or model_name in super().list_models():
            return super().transcribe_audio(
                audio_uri=audio_uri, model_name=model_name, **kwargs
            )

        # Check that the model name is a valid transcription model:
        # Valid models: whisper-1, gpt-4o-transcribe, gpt-4o-mini-transcribe, gpt-4o-transcribe-diarize
        valid_transcription_models = ["whisper", "transcribe"]
        if not any(valid in model_name for valid in valid_transcription_models):
            raise APIError(f"Model {model_name} does not support audio transcription.")

        # Get the local path to the audio file:
        audio_file_path = uri_to_local_file_path(audio_uri)

        # Open the audio file and transcribe:
        with open(audio_file_path, "rb") as audio_file:
            transcription = self.client.audio.transcriptions.create(
                model=model_name, file=audio_file, response_format="text"
            )

        # Update kwargs with other parameters:
        kwargs.update({"model_name": model_name})

        # Call the callback manager:
        self.callback_manager.on_audio_transcription(
            audio_uri=audio_uri, transcription=transcription, **kwargs
        )

        return transcription

    def generate_audio(
        self,
        text: str,
        voice: Optional[str] = None,
        audio_format: Optional[str] = None,
        model_name: Optional[str] = None,
        **kwargs,
    ) -> str:
        """Generate audio using OpenAI's TTS API."""

        # If no model name is provided, use the best model:
        if model_name is None:
            model_name = self.get_best_model(features=ModelFeatures.AudioGeneration)

        # If no model is available then delegate to superclass:
        if model_name is None or model_name in super().list_models():
            return super().generate_audio(
                text=text,
                voice=voice,
                audio_format=audio_format,
                model_name=model_name,
                **kwargs,
            )

        # If no voice is provided, use the default voice:
        if voice is None:
            voice = "onyx"

        # If no file format is provided, use mp3:
        if audio_format is None:
            audio_format = "mp3"

        # Call the OpenAI API to generate audio:
        response = self.client.audio.speech.create(
            model=model_name,
            voice=voice,
            response_format=audio_format,
            input=text,
            **kwargs,
        )

        # Save to temporary file:
        with tempfile.NamedTemporaryFile(
            suffix=f".{audio_format}", delete=False
        ) as temp_file:
            from litemind.utils.temp_file_manager import register_temp_file

            register_temp_file(temp_file.name)
            response.write_to_file(temp_file.name)
            audio_file_uri = "file://" + temp_file.name

            # Update kwargs with other parameters:
            kwargs.update(
                {"voice": voice, "audio_format": audio_format, "model_name": model_name}
            )

            # Call the callback manager:
            self.callback_manager.on_audio_generation(
                text=text, audio_uri=audio_file_uri, **kwargs
            )

            return audio_file_uri

    def generate_image(
        self,
        positive_prompt: str,
        negative_prompt: Optional[str] = None,
        model_name: str = None,
        image_width: int = 512,
        image_height: int = 512,
        preserve_aspect_ratio: bool = True,
        allow_resizing: bool = True,
        **kwargs,
    ) -> "Image":
        """Generate image using OpenAI's image generation API."""

        # Make prompt from positive and negative prompts:
        if negative_prompt:
            prompt = f"Image must consist in: {positive_prompt}\nBut NOT of: {negative_prompt}"
        else:
            prompt = positive_prompt

        # Get the best model if not provided:
        if model_name is None:
            model_name = self.get_best_model(features=ModelFeatures.ImageGeneration)

        # For Response API compatible models, try to use the newer approach
        if model_name and model_name in self._model_list:
            # Use the Response API for image generation if available
            try:
                response = self.client.responses.create(
                    model=model_name,
                    input=f"Generate an image: {prompt}",
                    tools=(
                        [{"type": "image_generation"}]
                        if "gpt-image" in model_name
                        else None
                    ),
                    **kwargs,
                )

                # Extract image from response - this would need to be implemented
                # based on how the Response API returns generated images
                # For now, fall back to the standard image generation API
                pass
            except Exception:
                pass

        # Fall back to standard DALL-E image generation API
        # (Same implementation as OpenAIApi)

        # Allowed sizes by OpenAI:
        if "dall-e-2" in model_name:
            allowed_sizes = [
                "256x256",
                "512x512",
                "1024x1024",
                "1792x1024",
                "1024x1792",
            ]
        elif "dall-e-3" in model_name:
            allowed_sizes = ["1024x1024", "1024x1792", "1792x1024"]
        elif "gpt-image" in model_name:
            # gpt-image models also support "auto" which lets the API choose
            allowed_sizes = ["1024x1024", "1024x1536", "1536x1024", "auto"]
        else:
            raise ValueError(
                f"Model {model_name} is not supported for image generation."
            )

        # Find appropriate size
        requested_size = f"{image_width}x{image_height}"
        if requested_size not in allowed_sizes:
            if not allow_resizing:
                raise ValueError(
                    f"Requested resolution {requested_size} is not allowed."
                )

            # Filter out non-dimensional sizes like "auto" before parsing
            dimensional_sizes = [s for s in allowed_sizes if "x" in s]

            allowed_widths = sorted(
                set(int(size.split("x")[0]) for size in dimensional_sizes)
            )
            allowed_heights = sorted(
                set(int(size.split("x")[1]) for size in dimensional_sizes)
            )

            closest_width = next(
                (w for w in allowed_widths if w >= image_width), allowed_widths[-1]
            )
            closest_height = next(
                (h for h in allowed_heights if h >= image_height), allowed_heights[-1]
            )
            closest_size = f"{closest_width}x{closest_height}"
        else:
            closest_size = requested_size

        # Set quality and format based on model:
        if "gpt-image" in model_name:
            # gpt-image models support quality: low, medium, high (default: high)
            if "quality" not in kwargs:
                kwargs["quality"] = "high"
            if "output_format" not in kwargs:
                kwargs["output_format"] = "png"
        elif "dall-e-3" in model_name:
            kwargs["quality"] = "hd"
            kwargs["response_format"] = "b64_json"
        elif "dall-e-2" in model_name:
            kwargs["response_format"] = "b64_json"

        # Generate image:
        response = self.client.images.generate(
            model=model_name,
            prompt=prompt,
            size=closest_size,
            n=1,
            **kwargs,
        )

        # Get the image and decode:
        image_b64 = response.data[0].b64_json
        if image_b64 is None:
            if hasattr(response.data[0], "url") and response.data[0].url:
                raise ValueError(
                    f"Model {model_name} returned a URL instead of base64 data. "
                    "Set response_format='b64_json' in kwargs."
                )
            raise ValueError(
                f"Model {model_name} did not return image data (b64_json is None)."
            )
        image_data = base64.b64decode(image_b64)

        # Create PIL image:
        from io import BytesIO

        from PIL import Image
        from PIL.Image import Resampling

        image = Image.open(BytesIO(image_data))

        # Resize if needed:
        if closest_size != requested_size:
            if preserve_aspect_ratio:
                image.thumbnail((image_width, image_height), Resampling.LANCZOS)
            else:
                image = image.resize((image_width, image_height), Resampling.LANCZOS)

        # Update kwargs and call callback:
        kwargs.update(
            {
                "model_name": model_name,
                "negative_prompt": negative_prompt,
                "image_width": image_width,
                "image_height": image_height,
                "preserve_aspect_ratio": preserve_aspect_ratio,
                "allow_resizing": allow_resizing,
            }
        )

        self.callback_manager.on_image_generation(image=image, prompt=prompt, **kwargs)
        return image

    def embed_texts(
        self,
        texts: Sequence[str],
        model_name: Optional[str] = None,
        dimensions: int = 512,
        **kwargs,
    ) -> Sequence[Sequence[float]]:
        """Generate text embeddings using OpenAI's embeddings API."""

        # Get the best model if not provided:
        if model_name is None:
            model_name = self.get_best_model(features=ModelFeatures.TextEmbeddings)

        # If the model belongs to the superclass, delegate:
        if model_name in super().list_models():
            return super().embed_texts(
                texts=texts, model_name=model_name, dimensions=dimensions, **kwargs
            )

        # Call OpenAI API to generate embeddings:
        response = self.client.embeddings.create(
            model=model_name, dimensions=dimensions, input=texts, **kwargs
        )

        # Extract embeddings from the response:
        embeddings = [e.embedding for e in response.data]

        # Update kwargs and call callback:
        kwargs.update({"model_name": model_name, "dimensions": dimensions})
        self.callback_manager.on_text_embedding(
            texts=texts, embeddings=embeddings, **kwargs
        )

        return embeddings
