"""
Automatic Code Execution Handler for LiteLLM Skills

When `litellm_code_execution` tool is present, this handler automatically:
1. Makes the LLM call
2. Executes any code the model generates
3. Continues the conversation with results
4. Returns final response with generated files inline (base64)

This mimics Anthropic's behavior where code execution happens automatically.
Generated files are returned directly in the response - no separate storage needed.
"""

import base64
import json
from enum import Enum
from typing import Any, Dict, List, Optional

from litellm._logging import verbose_logger


class LiteLLMInternalTools(str, Enum):
    """
    Enum for internal LiteLLM tools that are injected into requests.
    
    These tools are handled automatically by LiteLLM hooks and are not
    passed to the underlying LLM provider directly.
    """
    CODE_EXECUTION = "litellm_code_execution"


def get_litellm_code_execution_tool() -> Dict[str, Any]:
    """
    Returns the litellm_code_execution tool definition in OpenAI format.
    
    This tool enables automatic code execution in a sandboxed environment
    when skills include executable Python code.
    """
    return {
        "type": "function",
        "function": {
            "name": LiteLLMInternalTools.CODE_EXECUTION.value,
            "description": "Execute Python code in a sandboxed environment. Use this to run code that generates files, processes data, or performs computations. Generated files will be returned directly.",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "Python code to execute"
                    }
                },
                "required": ["code"]
            }
        }
    }


def get_litellm_code_execution_tool_anthropic() -> Dict[str, Any]:
    """
    Returns the litellm_code_execution tool definition in Anthropic/messages API format.
    
    This tool enables automatic code execution in a sandboxed environment
    when skills include executable Python code.
    """
    return {
        "name": LiteLLMInternalTools.CODE_EXECUTION.value,
        "description": "Execute Python code in a sandboxed environment. Use this to run code that generates files, processes data, or performs computations. Generated files will be returned directly.",
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Python code to execute"
                }
            },
            "required": ["code"]
        }
    }


# Singleton tool definition for backwards compatibility
LITELLM_CODE_EXECUTION_TOOL = get_litellm_code_execution_tool()


class CodeExecutionHandler:
    """
    Handles automatic code execution for LiteLLM skills.
    
    When enabled, this handler intercepts LLM responses with code execution
    tool calls, executes them in a sandbox, and continues the conversation
    automatically until completion.
    """
    
    def __init__(
        self,
        max_iterations: Optional[int] = None,
        sandbox_timeout: Optional[int] = None,
    ):
        from litellm.llms.litellm_proxy.skills.constants import (
            DEFAULT_MAX_ITERATIONS,
            DEFAULT_SANDBOX_TIMEOUT,
        )
        
        self.max_iterations = max_iterations or DEFAULT_MAX_ITERATIONS
        self.sandbox_timeout = sandbox_timeout or DEFAULT_SANDBOX_TIMEOUT
    
    async def execute_with_code_execution(
        self,
        model: str,
        messages: List[Dict],
        tools: List[Dict],
        skill_files: Dict[str, bytes],
        skill_id: Optional[str] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Execute an LLM call with automatic code execution handling.
        
        This method:
        1. Makes the initial LLM call
        2. If model calls litellm_code_execution, executes the code
        3. Continues conversation with results
        4. Repeats until model stops calling tools
        5. Returns final response with generated files inline
        
        Args:
            model: Model to use
            messages: Initial messages
            tools: Tools including litellm_code_execution
            skill_files: Dict of skill files for execution
            skill_id: Optional skill ID for tracking
            **kwargs: Additional args for litellm.acompletion
            
        Returns:
            Dict with:
            - response: Final LLM response
            - files: List of generated files with content (base64)
            - execution_results: List of code execution results
        """
        import litellm
        from litellm.llms.litellm_proxy.skills.sandbox_executor import (
            SkillsSandboxExecutor,
        )
        
        current_messages = list(messages)
        generated_files: List[Dict[str, Any]] = []  # Files returned directly
        execution_results: List[Dict] = []
        
        executor = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
        response: Any = None  # Initialize to avoid possibly unbound error
        
        for iteration in range(self.max_iterations):
            verbose_logger.debug(
                f"CodeExecutionHandler: Iteration {iteration + 1}/{self.max_iterations}"
            )
            
            # Make LLM call
            response = await litellm.acompletion(
                model=model,
                messages=current_messages,
                tools=tools,
                **kwargs,
            )
            
            assistant_message = response.choices[0].message  # type: ignore
            stop_reason = response.choices[0].finish_reason  # type: ignore
            
            # Build assistant message for conversation history
            assistant_msg_dict: Dict[str, Any] = {
                "role": "assistant",
                "content": assistant_message.content,
            }
            if assistant_message.tool_calls:
                assistant_msg_dict["tool_calls"] = [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments
                        }
                    }
                    for tc in assistant_message.tool_calls
                ]
            current_messages.append(assistant_msg_dict)
            
            # Check if we're done (no tool calls or not tool_calls finish reason)
            if stop_reason != "tool_calls" or not assistant_message.tool_calls:
                verbose_logger.debug(
                    f"CodeExecutionHandler: Completed after {iteration + 1} iterations"
                )
                return {
                    "response": response,
                    "files": generated_files,  # Files returned directly with base64 content
                    "execution_results": execution_results,
                    "messages": current_messages,
                }
            
            # Handle tool calls
            for tool_call in assistant_message.tool_calls:
                tool_name = tool_call.function.name
                
                if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value:
                    # Execute code in sandbox
                    try:
                        args = json.loads(tool_call.function.arguments)
                        code = args.get("code", "")
                        
                        verbose_logger.debug(
                            f"CodeExecutionHandler: Executing code ({len(code)} chars)"
                        )
                        
                        exec_result = executor.execute(
                            code=code,
                            skill_files=skill_files,
                        )

                        verbose_logger.debug(
                            f"CodeExecutionHandler: Execution result: {exec_result}"
                        )
                        
                        execution_results.append({
                            "iteration": iteration,
                            "success": exec_result["success"],
                            "output": exec_result["output"],
                            "error": exec_result["error"],
                            "files": [f["name"] for f in exec_result["files"]],
                        })
                        
                        # Build tool result content
                        tool_result = exec_result["output"] or ""
                        
                        # Collect generated files (returned directly, no storage)
                        if exec_result["files"]:
                            tool_result += "\n\nGenerated files:"
                            for f in exec_result["files"]:
                                file_content = base64.b64decode(f["content_base64"])
                                # Add to generated files list (returned in response)
                                generated_files.append({
                                    "name": f["name"],
                                    "mime_type": f["mime_type"],
                                    "content_base64": f["content_base64"],
                                    "size": len(file_content),
                                })
                                tool_result += f"\n- {f['name']} ({len(file_content)} bytes)"
                                
                                verbose_logger.debug(
                                    f"CodeExecutionHandler: Generated file {f['name']} ({len(file_content)} bytes)"
                                )
                        
                        if exec_result["error"]:
                            tool_result += f"\n\nError:\n{exec_result['error']}"
                        
                    except Exception as e:
                        tool_result = f"Code execution failed: {str(e)}"
                        execution_results.append({
                            "iteration": iteration,
                            "success": False,
                            "error": str(e),
                        })
                    
                    # Add tool result to messages
                    current_messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": tool_result,
                    })
                else:
                    # Non-code-execution tool - pass through
                    # In a full implementation, this would call other tool handlers
                    current_messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": f"Tool '{tool_name}' not handled by code execution handler",
                    })
        
        # Max iterations reached
        verbose_logger.warning(
            f"CodeExecutionHandler: Max iterations ({self.max_iterations}) reached"
        )
        return {
            "response": response,
            "files": generated_files,
            "execution_results": execution_results,
            "messages": current_messages,
            "max_iterations_reached": True,
        }


def has_code_execution_tool(tools: Optional[List[Dict]]) -> bool:
    """Check if litellm_code_execution tool is in the tools list."""
    if not tools:
        return False
    for tool in tools:
        func = tool.get("function", {})
        if func.get("name") == LiteLLMInternalTools.CODE_EXECUTION.value:
            return True
    return False


def add_code_execution_tool(tools: Optional[List[Dict]]) -> List[Dict]:
    """Add litellm_code_execution tool if not already present."""
    tools = tools or []
    if not has_code_execution_tool(tools):
        tools.append(LITELLM_CODE_EXECUTION_TOOL)
    return tools


# Global handler instance
code_execution_handler = CodeExecutionHandler()