from typing import Any, Callable
from functools import wraps
import inspect

import smolagents.default_tools
from smolagents._function_type_hints_utils import get_json_schema


SMOLAGENT_DEFAULT_TOOLS_OLLAMA_MAPPING = {
    smolagent_tool.name: {
        "type": "function",
        "function": {
            "name": smolagent_tool.name,
            "description": smolagent_tool.description,
            "parameters": {"type": "object", "properties": smolagent_tool.inputs},
            "required": list(smolagent_tool.inputs.keys()),
        },
    }
    for smolagent_tool in [
        class_element
        for _, class_element in inspect.getmembers(
            smolagents.default_tools, inspect.isclass
        )
        if ("Tool" in str(class_element)) and (getattr(class_element, "name", None))
    ]
}


class ToolsRegistry:
    _tools: dict[str, Any] = SMOLAGENT_DEFAULT_TOOLS_OLLAMA_MAPPING

    @classmethod
    def register(cls):
        def decorator(func: Callable):
            cls._tools.update({func.__name__: get_json_schema(func)})

            @wraps(func)
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)

            return wrapper

        return decorator

    @classmethod
    def get_all_tools(cls) -> list[dict[str, Any]]:
        return cls._tools