# createsonline/agents/tools.py
"""
Tool system for CREATESONLINE agents.

Provides a registry of callable tools/functions that agents can invoke
to interact with the outside world — databases, APIs, file systems, etc.
"""

import json
import time
import asyncio
import inspect
import traceback
from typing import Any, Callable, Dict, List, Optional, Union


class ToolResult:
    """Result of a tool invocation."""
    
    __slots__ = ('success', 'output', 'error', 'tool_name', 'duration', 'metadata')
    
    def __init__(self, success: bool, output: Any = None, error: str = None,
                 tool_name: str = '', duration: float = 0.0, metadata: dict = None):
        self.success = success
        self.output = output
        self.error = error
        self.tool_name = tool_name
        self.duration = duration
        self.metadata = metadata or {}
    
    def to_dict(self) -> dict:
        return {
            'success': self.success,
            'output': str(self.output) if self.output is not None else None,
            'error': self.error,
            'tool_name': self.tool_name,
            'duration': self.duration,
            'metadata': self.metadata,
        }
    
    def __repr__(self):
        status = 'OK' if self.success else 'FAIL'
        return f"ToolResult({status}, tool={self.tool_name!r})"


class Tool:
    """
    Wraps a callable function as an agent-invocable tool.
    
    Attributes:
        name: Unique identifier for the tool.
        description: Human-readable description (used in prompts).
        func: The actual callable (sync or async).
        parameters: JSON-schema-like dict describing accepted params.
        required_params: List of required parameter names.
        is_async: Whether the underlying function is async.
        tags: Optional tags for categorizing tools.
    """
    
    def __init__(self, name: str, func: Callable, description: str = '',
                 parameters: Dict[str, dict] = None, required_params: List[str] = None,
                 tags: List[str] = None):
        self.name = name
        self.func = func
        self.description = description or (func.__doc__ or '').strip()
        self.is_async = asyncio.iscoroutinefunction(func)
        self.tags = tags or []
        
        # Auto-infer parameters from function signature if not provided
        if parameters is not None:
            self.parameters = parameters
        else:
            self.parameters = self._infer_parameters(func)
        
        if required_params is not None:
            self.required_params = required_params
        else:
            self.required_params = self._infer_required(func)
    
    # ------------------------------------------------------------------
    @staticmethod
    def _infer_parameters(func: Callable) -> Dict[str, dict]:
        """Infer parameter schema from function signature."""
        sig = inspect.signature(func)
        params = {}
        type_map = {
            str: 'string', int: 'integer', float: 'number',
            bool: 'boolean', list: 'array', dict: 'object',
        }
        for pname, p in sig.parameters.items():
            if pname in ('self', 'cls'):
                continue
            ptype = 'string'
            if p.annotation != inspect.Parameter.empty:
                ptype = type_map.get(p.annotation, 'string')
            info: dict = {'type': ptype, 'description': ''}
            if p.default not in (inspect.Parameter.empty, None):
                info['default'] = p.default
            params[pname] = info
        return params
    
    @staticmethod
    def _infer_required(func: Callable) -> List[str]:
        sig = inspect.signature(func)
        required = []
        for pname, p in sig.parameters.items():
            if pname in ('self', 'cls'):
                continue
            if p.default is inspect.Parameter.empty:
                required.append(pname)
        return required
    
    # ------------------------------------------------------------------
    def validate_params(self, kwargs: dict) -> Optional[str]:
        """Return error message if params are invalid, else None."""
        for rp in self.required_params:
            if rp not in kwargs:
                return f"Missing required parameter: {rp}"
        for k in kwargs:
            if k not in self.parameters:
                return f"Unknown parameter: {k}"
        return None
    
    async def invoke(self, **kwargs) -> ToolResult:
        """Invoke the tool with the given keyword arguments."""
        err = self.validate_params(kwargs)
        if err:
            return ToolResult(success=False, error=err, tool_name=self.name)
        
        t0 = time.time()
        try:
            if self.is_async:
                output = await self.func(**kwargs)
            else:
                output = self.func(**kwargs)
            dt = time.time() - t0
            return ToolResult(success=True, output=output, tool_name=self.name, duration=dt)
        except Exception as exc:
            dt = time.time() - t0
            return ToolResult(
                success=False,
                error=f"{type(exc).__name__}: {exc}",
                tool_name=self.name,
                duration=dt,
                metadata={'traceback': traceback.format_exc()},
            )
    
    def to_schema(self) -> dict:
        """Return an OpenAI-compatible function calling schema."""
        return {
            'type': 'function',
            'function': {
                'name': self.name,
                'description': self.description,
                'parameters': {
                    'type': 'object',
                    'properties': self.parameters,
                    'required': self.required_params,
                },
            },
        }
    
    def __repr__(self):
        return f"Tool({self.name!r})"


class ToolRegistry:
    """
    Registry that holds tools available to agents.
    
    Usage::
    
        registry = ToolRegistry()
        
        @registry.tool("search_web", description="Search the web")
        def search_web(query: str) -> str:
            return requests.get(f"https://search.api?q={query}").text
        
        # Or register directly
        registry.register("calculate", calc_fn, "Evaluate math expression")
        
        # Invoke
        result = await registry.invoke("search_web", query="python docs")
    """
    
    def __init__(self):
        self._tools: Dict[str, Tool] = {}
        self._tag_index: Dict[str, List[str]] = {}   # tag -> [tool_names]
    
    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------
    def register(self, name: str, func: Callable, description: str = '',
                 parameters: Dict[str, dict] = None, required_params: List[str] = None,
                 tags: List[str] = None) -> Tool:
        """Register a function as a tool."""
        tool = Tool(name, func, description, parameters, required_params, tags)
        self._tools[name] = tool
        for tag in tool.tags:
            self._tag_index.setdefault(tag, []).append(name)
        return tool
    
    def tool(self, name: str = None, description: str = '', tags: List[str] = None):
        """Decorator to register a tool."""
        def decorator(func):
            tname = name or func.__name__
            self.register(tname, func, description, tags=tags)
            return func
        return decorator
    
    def unregister(self, name: str):
        """Remove a tool from the registry."""
        tool = self._tools.pop(name, None)
        if tool:
            for tag in tool.tags:
                lst = self._tag_index.get(tag, [])
                if name in lst:
                    lst.remove(name)
    
    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------
    def get(self, name: str) -> Optional[Tool]:
        return self._tools.get(name)
    
    def list_tools(self, tag: str = None) -> List[Tool]:
        """List all tools, optionally filtered by tag."""
        if tag:
            names = self._tag_index.get(tag, [])
            return [self._tools[n] for n in names if n in self._tools]
        return list(self._tools.values())
    
    def tool_names(self) -> List[str]:
        return list(self._tools.keys())
    
    def has_tool(self, name: str) -> bool:
        return name in self._tools
    
    # ------------------------------------------------------------------
    # Invocation
    # ------------------------------------------------------------------
    async def invoke(self, name: str, **kwargs) -> ToolResult:
        """Invoke a tool by name."""
        tool = self._tools.get(name)
        if tool is None:
            return ToolResult(success=False, error=f"Tool not found: {name}", tool_name=name)
        return await tool.invoke(**kwargs)
    
    def invoke_sync(self, name: str, **kwargs) -> ToolResult:
        """Synchronous wrapper for invoke."""
        loop = None
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            pass
        
        if loop and loop.is_running():
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as pool:
                future = pool.submit(asyncio.run, self.invoke(name, **kwargs))
                return future.result()
        else:
            return asyncio.run(self.invoke(name, **kwargs))
    
    # ------------------------------------------------------------------
    # Schema export (for LLM function calling)
    # ------------------------------------------------------------------
    def to_schema_list(self) -> List[dict]:
        """Return list of tool schemas for LLM function calling."""
        return [t.to_schema() for t in self._tools.values()]
    
    def to_prompt_description(self) -> str:
        """Return a text description of all tools for inclusion in prompts."""
        lines = ["Available tools:"]
        for t in self._tools.values():
            params_str = ', '.join(
                f"{k}: {v.get('type', 'any')}" for k, v in t.parameters.items()
            )
            lines.append(f"  - {t.name}({params_str}): {t.description}")
        return '\n'.join(lines)
    
    # ------------------------------------------------------------------
    # Built-in tools
    # ------------------------------------------------------------------
    def register_builtins(self):
        """Register a set of built-in utility tools."""
        
        @self.tool("echo", description="Echo the input text back", tags=["utility"])
        def echo(text: str) -> str:
            return text
        
        @self.tool("json_parse", description="Parse a JSON string", tags=["utility"])
        def json_parse(text: str) -> dict:
            return json.loads(text)
        
        @self.tool("json_format", description="Format data as JSON", tags=["utility"])
        def json_format(data: str, indent: int = 2) -> str:
            parsed = json.loads(data) if isinstance(data, str) else data
            return json.dumps(parsed, indent=indent, default=str)
        
        @self.tool("current_time", description="Get the current UTC time", tags=["utility"])
        def current_time() -> str:
            return time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime())
        
        @self.tool("calculate", description="Evaluate a safe math expression", tags=["math"])
        def calculate(expression: str) -> str:
            # Only allow safe math chars
            allowed = set('0123456789+-*/.() ')
            if not all(c in allowed for c in expression):
                return "Error: only numeric math expressions allowed"
            try:
                result = eval(expression, {"__builtins__": {}}, {})
                return str(result)
            except Exception as e:
                return f"Error: {e}"
    
    def __len__(self):
        return len(self._tools)
    
    def __contains__(self, name: str):
        return name in self._tools
    
    def __repr__(self):
        return f"ToolRegistry({len(self._tools)} tools)"
