from __future__ import annotations

"""
Response API backend implementation.
Standalone implementation optimized for the standard Response API format (originated by OpenAI).
"""

import os
from typing import Dict, List, Any, AsyncGenerator, Optional
from .base import LLMBackend, StreamChunk


class ResponseBackend(LLMBackend):
    """Backend using the standard Response API format."""

    def __init__(self, api_key: Optional[str] = None, **kwargs):
        super().__init__(api_key, **kwargs)
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")

    def convert_tools_to_response_api_format(
        self, tools: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Convert tools from Chat Completions format to Response API format if needed.

        Chat Completions format: {"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}}
        Response API format: {"type": "function", "name": ..., "description": ..., "parameters": ...}
        """
        if not tools:
            return tools

        converted_tools = []
        for tool in tools:
            if tool.get("type") == "function" and "function" in tool:
                # Chat Completions format - convert to Response API format
                func = tool["function"]
                converted_tools.append(
                    {
                        "type": "function",
                        "name": func["name"],
                        "description": func["description"],
                        "parameters": func.get("parameters", {}),
                    }
                )
            else:
                # Already in Response API format or non-function tool
                converted_tools.append(tool)

        return converted_tools

    def convert_messages_to_response_api_format(
        self, messages: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Convert messages from Chat Completions format to Response API format.

        Chat Completions tool message: {"role": "tool", "tool_call_id": "...", "content": "..."}
        Response API tool message: {"type": "function_call_output", "call_id": "...", "output": "..."}

        Note: Assistant messages with tool_calls should not be in input - they're generated by the backend.
        """
        converted_messages = []

        for message in messages:
            if message.get("role") == "tool":
                # Convert Chat Completions tool message to Response API format
                converted_messages.append(
                    {
                        "type": "function_call_output",
                        "call_id": message.get("tool_call_id"),
                        "output": message.get("content", ""),
                    }
                )
            elif message.get("type") == "function_call_output":
                # Already in Response API format - keep as-is
                converted_messages.append(message)
            elif message.get("role") == "assistant" and "tool_calls" in message:
                # Assistant message with tool_calls in native Responses API format
                # Remove tool_calls when sending as input - only results should be sent back
                cleaned_message = {
                    k: v for k, v in message.items() if k != "tool_calls"
                }
                converted_messages.append(cleaned_message)
            else:
                # Keep other message types as-is
                converted_messages.append(message)

        return converted_messages

    async def stream_with_tools(
        self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs
    ) -> AsyncGenerator[StreamChunk, None]:
        """Stream response using OpenAI Response API."""
        try:
            import openai

            client = openai.AsyncOpenAI(api_key=self.api_key)

            # Extract model and provider tool settings
            model = kwargs.get("model", "gpt-4o-mini")
            enable_web_search = kwargs.get("enable_web_search", False)
            enable_code_interpreter = kwargs.get("enable_code_interpreter", False)

            # Convert messages to Response API format (handles tool messages)
            converted_messages = self.convert_messages_to_response_api_format(messages)

            # Response API parameters (uses 'input', not 'messages')
            api_params = {"model": model, "input": converted_messages, "stream": True}

            # Add max_output_tokens if specified (o-series models don't support this)
            max_tokens = kwargs.get("max_tokens")
            if max_tokens and not model.startswith("o"):
                api_params["max_output_tokens"] = max_tokens

            # Add framework tools (convert to Response API format)
            if tools:
                converted_tools = self.convert_tools_to_response_api_format(tools)
                api_params["tools"] = converted_tools

            # Add provider tools (web search, code interpreter) if enabled
            provider_tools = []
            if enable_web_search:
                provider_tools.append({"type": "web_search"})

            if enable_code_interpreter:
                provider_tools.append(
                    {"type": "code_interpreter", "container": {"type": "auto"}}
                )

            if provider_tools:
                if "tools" not in api_params:
                    api_params["tools"] = []
                api_params["tools"].extend(provider_tools)

            stream = await client.responses.create(**api_params)

            content = ""

            async for chunk in stream:
                # Handle Responses API streaming format
                if hasattr(chunk, "type"):
                    if chunk.type == "response.output_text.delta" and hasattr(
                        chunk, "delta"
                    ):
                        content += chunk.delta
                        yield StreamChunk(type="content", content=chunk.delta)
                    elif chunk.type == "response.web_search_call.in_progress":
                        yield StreamChunk(
                            type="content",
                            content=f"\n🔍 [Provider Tool: Web Search] Starting search...",
                        )
                    elif chunk.type == "response.web_search_call.searching":
                        yield StreamChunk(
                            type="content",
                            content=f"🔍 [Provider Tool: Web Search] Searching...",
                        )
                    elif chunk.type == "response.web_search_call.completed":
                        yield StreamChunk(
                            type="content",
                            content=f"✅ [Provider Tool: Web Search] Search completed",
                        )
                    elif chunk.type == "response.code_interpreter_call.in_progress":
                        yield StreamChunk(
                            type="content",
                            content=f"\n💻 [Provider Tool: Code Interpreter] Starting execution...",
                        )
                    elif chunk.type == "response.code_interpreter_call.executing":
                        yield StreamChunk(
                            type="content",
                            content=f"💻 [Provider Tool: Code Interpreter] Executing...",
                        )
                    elif chunk.type == "response.code_interpreter_call.completed":
                        yield StreamChunk(
                            type="content",
                            content=f"✅ [Provider Tool: Code Interpreter] Execution completed",
                        )
                    elif chunk.type == "response.output_item.done":
                        # Get search query or executed code details - show them right after completion
                        if hasattr(chunk, "item") and chunk.item:
                            if (
                                hasattr(chunk.item, "type")
                                and chunk.item.type == "web_search_call"
                            ):
                                if hasattr(chunk.item, "action") and hasattr(
                                    chunk.item.action, "query"
                                ):
                                    search_query = chunk.item.action.query
                                    if search_query:
                                        yield StreamChunk(
                                            type="content",
                                            content=f"🔍 [Search Query] '{search_query}'",
                                        )
                            elif (
                                hasattr(chunk.item, "type")
                                and chunk.item.type == "code_interpreter_call"
                            ):
                                if hasattr(chunk.item, "code") and chunk.item.code:
                                    # Format code as a proper code block - don't assume language
                                    yield StreamChunk(
                                        type="content",
                                        content=f"💻 [Code Executed]\n```\n{chunk.item.code}\n```\n",
                                    )

                                # Also show the execution output if available
                                if (
                                    hasattr(chunk.item, "outputs")
                                    and chunk.item.outputs
                                ):
                                    for output in chunk.item.outputs:
                                        output_text = None
                                        if hasattr(output, "text") and output.text:
                                            output_text = output.text
                                        elif (
                                            hasattr(output, "content")
                                            and output.content
                                        ):
                                            output_text = output.content
                                        elif hasattr(output, "data") and output.data:
                                            output_text = str(output.data)
                                        elif isinstance(output, str):
                                            output_text = output
                                        elif isinstance(output, dict):
                                            # Handle dict format outputs
                                            if "text" in output:
                                                output_text = output["text"]
                                            elif "content" in output:
                                                output_text = output["content"]
                                            elif "data" in output:
                                                output_text = str(output["data"])

                                        if output_text and output_text.strip():
                                            yield StreamChunk(
                                                type="content",
                                                content=f"📊 [Result] {output_text.strip()}\n",
                                            )
                    elif chunk.type == "response.completed":
                        # Extract and yield tool calls from the complete response
                        if hasattr(chunk, "response"):
                            response_dict = self._convert_to_dict(chunk.response)

                            # Extract builtin tool results from output array
                            builtin_tool_results = []
                            if (
                                isinstance(response_dict, dict)
                                and "output" in response_dict
                            ):
                                for item in response_dict["output"]:
                                    if item.get("type") == "code_interpreter_call":
                                        # Code execution result
                                        builtin_tool_results.append(
                                            {
                                                "id": item.get("id", ""),
                                                "tool_type": "code_interpreter",
                                                "status": item.get("status"),
                                                "code": item.get("code", ""),
                                                "outputs": item.get("outputs"),
                                                "container_id": item.get(
                                                    "container_id"
                                                ),
                                            }
                                        )
                                    elif item.get("type") == "web_search_call":
                                        # Web search result
                                        builtin_tool_results.append(
                                            {
                                                "id": item.get("id", ""),
                                                "tool_type": "web_search",
                                                "status": item.get("status"),
                                                "query": item.get("query", ""),
                                                "results": item.get("results"),
                                            }
                                        )

                            # Yield builtin tool results if any were found
                            if builtin_tool_results:
                                yield StreamChunk(
                                    type="builtin_tool_results",
                                    builtin_tool_results=builtin_tool_results,
                                )

                            # Yield the complete response for internal use
                            yield StreamChunk(
                                type="complete_response", response=response_dict
                            )
                        else:
                            # Fallback if no response object
                            complete_message = {
                                "role": "assistant",
                                "content": content.strip(),
                            }
                            yield StreamChunk(
                                type="complete_message",
                                complete_message=complete_message,
                            )

                        # Signal completion
                        yield StreamChunk(type="done")

        except Exception as e:
            yield StreamChunk(type="error", error=str(e))

    def get_provider_name(self) -> str:
        """Get the provider name."""
        return "OpenAI"

    def get_supported_builtin_tools(self) -> List[str]:
        """Get list of builtin tools supported by OpenAI."""
        return ["web_search", "code_interpreter"]

    def extract_tool_name(self, tool_call: Dict[str, Any]) -> str:
        """Extract tool name from OpenAI format (handles both Chat Completions and Responses API)."""
        # Check if it's Chat Completions format
        if "function" in tool_call:
            return tool_call.get("function", {}).get("name", "unknown")
        # Otherwise assume Responses API format
        return tool_call.get("name", "unknown")

    def extract_tool_arguments(self, tool_call: Dict[str, Any]) -> Dict[str, Any]:
        """Extract tool arguments from OpenAI format (handles both Chat Completions and Responses API)."""
        # Check if it's Chat Completions format
        if "function" in tool_call:
            return tool_call.get("function", {}).get("arguments", {})
        # Otherwise assume Responses API format
        arguments = tool_call.get("arguments", {})
        if isinstance(arguments, str):
            try:
                import json

                return json.loads(arguments)
            except:
                return {}
        return arguments

    def extract_tool_call_id(self, tool_call: Dict[str, Any]) -> str:
        """Extract tool call ID from OpenAI format (handles both Chat Completions and Responses API)."""
        # For Responses API, use call_id (for tool results), for Chat Completions use id
        return tool_call.get("call_id") or tool_call.get("id") or ""

    def create_tool_result_message(
        self, tool_call: Dict[str, Any], result_content: str
    ) -> Dict[str, Any]:
        """Create tool result message for OpenAI Responses API format."""
        tool_call_id = self.extract_tool_call_id(tool_call)
        # Use Responses API format directly - no conversion needed
        return {
            "type": "function_call_output",
            "call_id": tool_call_id,
            "output": result_content,
        }

    def extract_tool_result_content(self, tool_result_message: Dict[str, Any]) -> str:
        """Extract content from OpenAI Responses API tool result message."""
        return tool_result_message.get("output", "")

    def _convert_to_dict(self, obj) -> Dict[str, Any]:
        """Convert any object to dictionary with multiple fallback methods."""
        try:
            if hasattr(obj, "model_dump"):
                return obj.model_dump()
            elif hasattr(obj, "dict"):
                return obj.dict()
            else:
                return dict(obj)
        except:
            # Final fallback: extract key attributes manually
            return {
                key: getattr(obj, key, None)
                for key in dir(obj)
                if not key.startswith("_") and not callable(getattr(obj, key, None))
            }

    def estimate_tokens(self, text: str) -> int:
        """Estimate token count for text (rough approximation)."""
        return len(text) // 4

    def calculate_cost(
        self, input_tokens: int, output_tokens: int, model: str
    ) -> float:
        """Calculate cost for OpenAI token usage (2024-2025 pricing)."""
        model_lower = model.lower()

        if "gpt-4" in model_lower:
            if "4o-mini" in model_lower:
                input_cost = input_tokens * 0.00015 / 1000
                output_cost = output_tokens * 0.0006 / 1000
            elif "4o" in model_lower:
                input_cost = input_tokens * 0.005 / 1000
                output_cost = output_tokens * 0.020 / 1000
            else:
                input_cost = input_tokens * 0.03 / 1000
                output_cost = output_tokens * 0.06 / 1000
        elif "gpt-3.5" in model_lower:
            input_cost = input_tokens * 0.0005 / 1000
            output_cost = output_tokens * 0.0015 / 1000
        else:
            input_cost = input_tokens * 0.0005 / 1000
            output_cost = output_tokens * 0.0015 / 1000

        return input_cost + output_cost
