"""
Uni-Wizard Tool Registry
Central registry for all tool capabilities
"""

import json
import inspect
from typing import Dict, Callable, Any, Optional
from dataclasses import dataclass, asdict
from functools import wraps


@dataclass
class ToolSchema:
    """Schema definition for a tool"""
    name: str
    description: str
    parameters: Dict[str, Any]
    returns: str
    examples: list = None
    
    def to_dict(self):
        return asdict(self)


@dataclass
class ToolResult:
    """Standardized tool execution result"""
    success: bool
    data: Any
    error: Optional[str] = None
    execution_time_ms: Optional[float] = None
    
    def to_json(self) -> str:
        return json.dumps({
            'success': self.success,
            'data': self.data,
            'error': self.error,
            'execution_time_ms': self.execution_time_ms
        }, indent=2)
    
    def __str__(self) -> str:
        if self.success:
            return str(self.data)
        return f"Error: {self.error}"


class ToolRegistry:
    """
    Central registry for all uni-wizard tools.
    
    All tools register here with their schemas.
    The LLM queries available tools via get_tool_definitions().
    """
    
    def __init__(self):
        self._tools: Dict[str, Dict] = {}
        self._categories: Dict[str, list] = {}
    
    def register(
        self,
        name: str,
        handler: Callable,
        description: str = None,
        parameters: Dict = None,
        category: str = "general",
        examples: list = None
    ):
        """
        Register a tool in the registry.
        
        Args:
            name: Tool name (used in tool calls)
            handler: Function to execute
            description: What the tool does
            parameters: JSON Schema for parameters
            category: Tool category (system, git, network, file)
            examples: Example usages
        """
        # Auto-extract description from docstring if not provided
        if description is None and handler.__doc__:
            description = handler.__doc__.strip().split('\n')[0]
        
        # Auto-extract parameters from function signature
        if parameters is None:
            parameters = self._extract_params(handler)
        
        self._tools[name] = {
            'name': name,
            'handler': handler,
            'description': description or f"Execute {name}",
            'parameters': parameters,
            'category': category,
            'examples': examples or []
        }
        
        # Add to category
        if category not in self._categories:
            self._categories[category] = []
        self._categories[category].append(name)
        
        return self  # For chaining
    
    def _extract_params(self, handler: Callable) -> Dict:
        """Extract parameter schema from function signature"""
        sig = inspect.signature(handler)
        params = {
            "type": "object",
            "properties": {},
            "required": []
        }
        
        for name, param in sig.parameters.items():
            # Skip 'self', 'cls', and params with defaults
            if name in ('self', 'cls'):
                continue
            
            param_info = {"type": "string"}  # Default
            
            # Try to infer type from annotation
            if param.annotation != inspect.Parameter.empty:
                if param.annotation == int:
                    param_info["type"] = "integer"
                elif param.annotation == float:
                    param_info["type"] = "number"
                elif param.annotation == bool:
                    param_info["type"] = "boolean"
                elif param.annotation == list:
                    param_info["type"] = "array"
                elif param.annotation == dict:
                    param_info["type"] = "object"
            
            # Add description if in docstring
            if handler.__doc__:
                # Simple param extraction from docstring
                for line in handler.__doc__.split('\n'):
                    if f'{name}:' in line or f'{name} (' in line:
                        desc = line.split(':', 1)[-1].strip()
                        param_info["description"] = desc
                        break
            
            params["properties"][name] = param_info
            
            # Required if no default
            if param.default == inspect.Parameter.empty:
                params["required"].append(name)
        
        return params
    
    def execute(self, name: str, **params) -> ToolResult:
        """
        Execute a tool by name with parameters.
        
        Args:
            name: Tool name
            **params: Tool parameters
            
        Returns:
            ToolResult with success/failure and data
        """
        import time
        start = time.time()
        
        tool = self._tools.get(name)
        if not tool:
            return ToolResult(
                success=False,
                data=None,
                error=f"Tool '{name}' not found in registry",
                execution_time_ms=(time.time() - start) * 1000
            )
        
        try:
            handler = tool['handler']
            result = handler(**params)
            
            return ToolResult(
                success=True,
                data=result,
                execution_time_ms=(time.time() - start) * 1000
            )
            
        except Exception as e:
            return ToolResult(
                success=False,
                data=None,
                error=f"{type(e).__name__}: {str(e)}",
                execution_time_ms=(time.time() - start) * 1000
            )
    
    def get_tool(self, name: str) -> Optional[Dict]:
        """Get tool definition by name"""
        tool = self._tools.get(name)
        if tool:
            # Return without handler (not serializable)
            return {
                'name': tool['name'],
                'description': tool['description'],
                'parameters': tool['parameters'],
                'category': tool['category'],
                'examples': tool['examples']
            }
        return None
    
    def get_tools_by_category(self, category: str) -> list:
        """Get all tools in a category"""
        tool_names = self._categories.get(category, [])
        return [self.get_tool(name) for name in tool_names if self.get_tool(name)]
    
    def list_tools(self, category: str = None) -> list:
        """List all tool names, optionally filtered by category"""
        if category:
            return self._categories.get(category, [])
        return list(self._tools.keys())
    
    def get_tool_definitions(self) -> str:
        """
        Get all tool definitions formatted for LLM system prompt.
        Returns JSON string of all tools with schemas.
        """
        tools = []
        for name, tool in self._tools.items():
            tools.append({
                "name": name,
                "description": tool['description'],
                "parameters": tool['parameters']
            })
        
        return json.dumps(tools, indent=2)
    
    def get_categories(self) -> list:
        """Get all tool categories"""
        return list(self._categories.keys())


# Global registry instance
registry = ToolRegistry()


def tool(name: str = None, category: str = "general", examples: list = None):
    """
    Decorator to register a function as a tool.
    
    Usage:
        @tool(category="system")
        def system_info():
            return {...}
    """
    def decorator(func: Callable):
        tool_name = name or func.__name__
        registry.register(
            name=tool_name,
            handler=func,
            category=category,
            examples=examples
        )
        return func
    return decorator


# Convenience function for quick tool execution
def call_tool(name: str, **params) -> str:
    """Execute a tool and return string result"""
    result = registry.execute(name, **params)
    return str(result)
