"""
UltraGPT Tools Module

This module contains all tool-related functionality including:
- Internal tool loading and execution
- User tool validation and conversion
- Native tool format conversions
- Tool execution orchestration

All these functions were moved from core.py to keep the main class organized.
"""

import os
import json
import inspect
import importlib
from pydantic import BaseModel

from . import config

class ToolManager:
    """
    Manages all tool-related operations for UltraGPT.
    This class handles both internal tools and user-defined tools.
    """
    
    def __init__(self, ultragpt_instance):
        """
        Initialize the ToolManager with a reference to the UltraGPT instance.
        
        Args:
            ultragpt_instance: The UltraGPT instance that owns this tool manager
        """
        self.ultragpt = ultragpt_instance
        self.log = ultragpt_instance.log
        self.verbose = ultragpt_instance.verbose
        self.provider_manager = ultragpt_instance.provider_manager
        self.max_tokens = ultragpt_instance.max_tokens
        self.google_api_key = ultragpt_instance.google_api_key
        self.search_engine_id = ultragpt_instance.search_engine_id

    def load_internal_tools(self, tools: list) -> dict:
        """Dynamically load internal tool configurations from tool directories"""
        tools_base_path = os.path.join(os.path.dirname(__file__), 'tools')
        loaded_tools = {}
        
        for tool_name in tools:
            tool_name_normalized = tool_name.replace('-', '_')  # Convert web-search to web_search
            tool_path = os.path.join(tools_base_path, tool_name_normalized)
            
            if not os.path.exists(tool_path):
                self.log.warning(f"Tool directory not found: {tool_path}")
                continue
                
            try:
                # Load prompts module to get _info and _description
                prompts_module = importlib.import_module(f'.tools.{tool_name_normalized}.prompts', package='ultragpt')
                info = getattr(prompts_module, '_info', f"Tool: {tool_name}")
                description = getattr(prompts_module, '_description', info)
                
                # Load schemas module to get schema classes
                schemas_module = importlib.import_module(f'.tools.{tool_name_normalized}.schemas', package='ultragpt')
                
                # Find the main schema class (usually ends with 'Query')
                schema_class = None
                for name, obj in inspect.getmembers(schemas_module):
                    if (inspect.isclass(obj) and 
                        issubclass(obj, BaseModel) and 
                        name.endswith('Query')):
                        schema_class = obj
                        break
                
                if not schema_class:
                    self.log.warning(f"No Query schema found for tool: {tool_name}")
                    continue
                
                # Convert schema to tool format for native calling
                schema_dict = schema_class.model_json_schema()
                
                loaded_tools[tool_name_normalized] = {  # Use normalized name as key
                    'name': tool_name_normalized,
                    'display_name': tool_name,
                    'description': description,
                    'schema': schema_class,
                    'native_schema': {
                        'type': 'function',
                        'function': {
                            'name': tool_name_normalized,
                            'description': description,
                            'parameters': schema_dict
                        }
                    }
                }
                
                self.log.debug(f"✓ Loaded tool: {tool_name} (as {tool_name_normalized})")
                
            except Exception as e:
                self.log.warning(f"Failed to load tool {tool_name}: {str(e)}")
                continue
        
        return loaded_tools

    def convert_internal_tools_to_native_format(self, loaded_tools: dict) -> list:
        """Convert loaded internal tools to native AI provider tool format"""
        native_tools = []
        
        for tool_name, tool_config in loaded_tools.items():
            # Use the native_schema that was already created in load_internal_tools
            if 'native_schema' in tool_config:
                native_tools.append(tool_config['native_schema'])
            else:
                # Fallback: create schema from tool config
                native_tool = {
                    "type": "function",
                    "function": {
                        "name": tool_config['name'],
                        "description": tool_config['description'],
                        "parameters": {}
                    }
                }
                native_tools.append(native_tool)
        
        return native_tools

    def execute_internal_tool_with_params(
        self, 
        tool_config: dict, 
        parameters: dict, 
        tools_config: dict
    ) -> str:
        """Execute an internal tool directly with AI-provided parameters"""
        # but kept for future extensibility
        try:
            # Get tool-specific config
            tool_name_normalized = tool_config['name']  # Use normalized name
            config_data = tools_config.get(tool_name_normalized, {})
            
            # Import and execute the tool's execute_tool function
            try:
                core_module = importlib.import_module(f'.tools.{tool_name_normalized}.core', package='ultragpt')
                execute_function = getattr(core_module, 'execute_tool', None)
                
                if execute_function:
                    # Set credentials for web search tool in thread-local context
                    if tool_name_normalized == "web_search":
                        try:
                            context_module = importlib.import_module(f'.tools.{tool_name_normalized}.context', package='ultragpt')
                            context_module.set_credentials(self.google_api_key, self.search_engine_id)
                        except ImportError:
                            pass  # Context module not available
                    
                    result = execute_function(parameters)
                    
                    # Clear credentials after execution
                    if tool_name_normalized == "web_search":
                        try:
                            context_module.clear_credentials()
                        except:
                            pass
                    
                    return result if result else ""
                else:
                    return f"No execute_tool function found for {tool_name_normalized}"
                    
            except Exception as e:
                return f"Error executing tool {tool_name_normalized}: {str(e)}"
                
        except Exception as e:
            return f"Tool execution error: {str(e)}"

    def execute_tools(
        self,
        history: list,
        tools: list,
        tools_config: dict
    ) -> tuple:
        """Execute tools using native AI tool calling - returns (result_string, tool_usage_details)"""
        tool_usage_details = []
        
        if not tools or not history:
            return "", tool_usage_details
        
        # Get the latest user message from history
        message = ""
        for msg in reversed(history):
            if msg.get("role") == "user" and msg.get("content"):
                message = msg["content"]
                break
        
        if not message:
            return "", tool_usage_details
        
        try:
            self.log.info(f"Loading and executing {len(tools)} tools using native AI tool calling")
            if self.verbose:
                self.log.debug(f"➤ Loading {len(tools)} tools for native AI tool calling")
                self.log.debug(f"Query: {message[:100] + '...' if len(message) > 100 else message}")
                self.log.debug("-" * 40)
            
            # Load internal tools dynamically
            loaded_tools = self.load_internal_tools(tools)
            if not loaded_tools:
                self.log.warning("No tools could be loaded")
                return "", tool_usage_details
            
            # Convert to native tool format
            native_tools = self.convert_internal_tools_to_native_format(loaded_tools)
            
            # Create tool selection prompt
            tool_descriptions = []
            for tool_name, tool_config in loaded_tools.items():
                tool_descriptions.append(f"- {tool_config['display_name']}: {tool_config['description']}")
            
            tool_selection_prompt = f"""
Available internal tools:
{chr(10).join(tool_descriptions)}

Analyze the user's message and select the appropriate tools with their parameters. Each tool should be called with the specific parameters needed to help answer the user's question.

User message: "{message}"

IMPORTANT: 
- Only call tools that are actually needed to help answer the user's question
- Use the tool's schema to provide the correct parameters
- If no tools are needed, don't call any tools
- You can call multiple tools if needed
"""
            
            # Prepare messages for tool calling
            tool_messages = [
                {"role": "system", "content": tool_selection_prompt},
                {"role": "user", "content": message}
            ]
            
            # Add conversation history context (limited)
            if history:
                context_messages = history[-config.MAX_CONTEXT_MESSAGES:]  # Last N messages for context
                for msg in context_messages:
                    if msg.get("content") and msg.get("role") in ["user", "assistant"]:
                        tool_messages.insert(-1, {
                            "role": msg["role"], 
                            "content": msg["content"]
                        })
            
            # Get model from tools_config or use default
            model = config.DEFAULT_TOOLS_MODEL
            parallel_tool_calls = True  # Default to enabling parallel tool calls
            for tool_config in tools_config.values():
                if 'model' in tool_config:
                    model = tool_config['model']
                if 'parallel_tool_calls' in tool_config:
                    parallel_tool_calls = tool_config['parallel_tool_calls']
                if 'model' in tool_config or 'parallel_tool_calls' in tool_config:
                    break
            
            # Make native tool call
            try:
                response_message, tokens = self.provider_manager.chat_completion_with_tools(
                    model=model,
                    messages=tool_messages,
                    tools=native_tools,
                    temperature=config.TOOL_SELECTION_TEMPERATURE,  # Low temperature for tool selection
                    max_tokens=self.max_tokens,
                    parallel_tool_calls=parallel_tool_calls  # Enable parallel tool calls
                )
                
                if self.verbose:
                    self.log.debug(f"AI tool selection completed (tokens: {tokens})")
                
            except Exception as e:
                # If native tool calling fails, fall back to no tools
                self.log.warning(f"Native tool calling failed, proceeding without tools: {str(e)}")
                return ""
            
            # Process tool calls if any were made
            if not response_message.get('tool_calls'):
                if self.verbose:
                    self.log.debug("AI decided no tools are needed")
                return "", tool_usage_details
            
            # Execute the selected tools
            tool_results = []
            for tool_call in response_message.get('tool_calls', []):
                function_name = tool_call.get('function', {}).get('name')
                function_args = tool_call.get('function', {}).get('arguments', {})
                
                # Parse arguments if they're in string format
                if isinstance(function_args, str):
                    try:
                        function_args = json.loads(function_args)
                    except json.JSONDecodeError:
                        self.log.error(f"Failed to parse tool arguments: {function_args}")
                        continue
                
                if function_name in loaded_tools:
                    tool_config = loaded_tools[function_name]
                    
                    try:
                        if self.verbose:
                            self.log.debug(f"Executing tool: {function_name}")
                            self.log.debug(f"Parameters: {function_args}")
                        
                        # Execute the tool with the AI-selected parameters
                        tool_result = self.execute_internal_tool_with_params(
                            tool_config, function_args, tools_config
                        )
                        
                        tool_results.append({
                            "tool": tool_config['display_name'],
                            "response": tool_result
                        })
                        
                        # Track tool usage
                        tool_usage_details.append({
                            "tool_name": function_name,
                            "display_name": tool_config['display_name'],
                            "parameters": function_args,
                            "result": tool_result,
                            "success": True,
                            "error": None
                        })
                        
                        if self.verbose:
                            self.log.debug(f"✓ {function_name} completed")
                            self.log.debug("-" * 40)
                            self.log.debug(tool_result if tool_result else "(empty result)")
                            self.log.debug("-" * 40)
                            
                    except Exception as e:
                        error_msg = f"Tool execution failed: {str(e)}"
                        self.log.error(f"Tool {function_name} execution failed: {str(e)}")
                        tool_results.append({
                            "tool": tool_config['display_name'],
                            "response": error_msg
                        })
                        
                        # Track failed tool usage
                        tool_usage_details.append({
                            "tool_name": function_name,
                            "display_name": tool_config['display_name'],
                            "parameters": function_args,
                            "result": None,
                            "success": False,
                            "error": str(e)
                        })
                else:
                    self.log.warning(f"Unknown tool called: {function_name}")
            
            # Format results
            if not tool_results:
                return "", tool_usage_details
                
            formatted_responses = []
            for result in tool_results:
                tool_name = result['tool'].upper()
                response = result['response'].strip() if result['response'] else ""
                if response:
                    formatted = f"[{tool_name}]\n{response}"
                    formatted_responses.append(formatted)
            
            success_count = len([r for r in tool_results if r['response'] and not r['response'].startswith('Tool execution failed')])
            self.log.info(f"Tools execution completed ({success_count}/{len(tool_results)} successful)")
            
            if self.verbose:
                self.log.debug(f"✓ Tools execution completed ({success_count}/{len(tool_results)} successful)")
                
            return "\n\n".join(formatted_responses), tool_usage_details
                
        except Exception as e:
            self.log.error(f"Tool execution failed: {str(e)}")
            if self.verbose:
                self.log.debug(f"✗ Tool execution failed: {str(e)}")
            return "", tool_usage_details

    def convert_user_tools_to_native_format(self, user_tools: list) -> list:
        """Convert UserTool objects to native AI provider tool format"""
        native_tools = []
        
        for tool in user_tools:
            if isinstance(tool, dict):
                tool_dict = tool
            elif hasattr(tool, 'model_dump'):
                tool_dict = tool.model_dump()
            else:
                self.log.warning("Invalid tool format: " + str(type(tool)))
                continue
            
            # Get parameters schema and ensure it has additionalProperties: false for OpenAI strict mode
            parameters_schema = tool_dict["parameters_schema"].copy()
            
            # Surgically add reasoning and stop_after_tool_call parameters to the schema
            if "properties" not in parameters_schema:
                parameters_schema["properties"] = {}
            
            # Add reasoning parameter
            parameters_schema["properties"]["reasoning"] = {
                "type": "string",
                "description": "Detailed reasoning for why this tool was chosen and how it will help solve the user's request"
            }
            
            # Add stop_after_tool_call parameter  
            parameters_schema["properties"]["stop_after_tool_call"] = {
                "type": "boolean",
                "description": "Whether to stop execution after this tool call (true if task is complete or user input needed, false to continue with more tools)"
            }
            
            # Ensure additionalProperties is false and required includes all properties for OpenAI strict mode
            def ensure_openai_strict_compliance(schema):
                if isinstance(schema, dict):
                    if schema.get("type") == "object":
                        schema["additionalProperties"] = False
                        # For OpenAI strict mode, required must include ALL properties if any are specified
                        if "properties" in schema:
                            all_properties = list(schema["properties"].keys())
                            schema["required"] = all_properties
                    
                    for key, value in schema.items():
                        if key == "properties" and isinstance(value, dict):
                            for prop_value in value.values():
                                ensure_openai_strict_compliance(prop_value)
                        elif isinstance(value, dict):
                            ensure_openai_strict_compliance(value)
                        elif isinstance(value, list):
                            for item in value:
                                ensure_openai_strict_compliance(item)
            
            ensure_openai_strict_compliance(parameters_schema)
            
            # Convert to OpenAI function calling format (Claude will handle conversion)
            native_tool = {
                "type": "function",
                "function": {
                    "name": tool_dict["name"],
                    "description": tool_dict["description"],
                    "parameters": parameters_schema,
                    "strict": True
                }
            }
            native_tools.append(native_tool)
        
        return native_tools

    def validate_user_tools(self, user_tools: list) -> list:
        """Validate and format user tools (both UserTool and ExpertTool)"""
        validated_tools = []
        
        for tool in user_tools:
            if isinstance(tool, dict):
                # Ensure all required fields are present for UserTool
                required_fields = ['name', 'description', 'parameters_schema', 'usage_guide', 'when_to_use']
                if all(field in tool for field in required_fields):
                    validated_tools.append(tool)
                else:
                    missing = [field for field in required_fields if field not in tool]
                    self.log.warning("Tool missing required fields: " + str(missing))
                    if self.verbose:
                        self.log.debug("⚠ Tool missing fields: " + str(missing))
            elif hasattr(tool, 'model_dump'):
                # Pydantic model (UserTool or ExpertTool)
                validated_tools.append(tool.model_dump())
            else:
                self.log.warning("Invalid tool format: " + str(type(tool)))
                if self.verbose:
                    self.log.debug("⚠ Invalid tool format: " + str(type(tool)))
        
        return validated_tools
