"""Tool Registry - Single source of truth for tool management.

This module provides the centralized tool registry that all components use.

ALL PYDANTIC — ToolDefinition is a Pydantic model, parameters is JSONSchema,
to_schema() returns LLMToolSpec, call() takes ToolArgs.

Usage:
    from core.tools import ToolRegistry, get_tool_registry

    # Get global registry
    registry = get_tool_registry()

    # Register a tool
    registry.register(tool_def)

    # Get a tool
    tool = registry.get("web.search")

    # Call a tool
    result = await registry.call("web.search", ToolArgs(query="hello"), ctx)

    # Get schemas for patterns
    schemas = registry.get_schemas(["web.*", "workspace.*"])
"""

import fnmatch
import logging
import threading
from collections.abc import Awaitable, Callable

from pydantic import BaseModel, Field

from core.common import JsonObject
from core.pathways.models.common import JSONSchema
from core.pathways.models.tool import ToolArgs, ToolOutput
from core.types.llm import LLMFunctionParameters, LLMFunctionSpec, LLMToolSpec

logger = logging.getLogger(__name__)


# =============================================================================
# TYPES
# =============================================================================

# Canonical handler signature: async (ToolArgs, ctx) -> ToolOutput
# Pydantic end-to-end: no dict[str, Any] in the pipeline.
ToolHandler = Callable[[ToolArgs, object], Awaitable[ToolOutput]]


class ToolDefinition(BaseModel):
    """Definition of a tool — Pydantic model, not a dataclass.

    parameters is a JSONSchema model (not dict[str, Any]).
    to_schema() returns an LLMToolSpec model.
    """

    model_config = {"arbitrary_types_allowed": True}

    name: str = Field(description="e.g., 'web.search'")
    description: str = ""
    parameters: JSONSchema | JsonObject = Field(default_factory=lambda: JSONSchema(properties={}))
    handler: ToolHandler | None = Field(default=None, exclude=True)

    # Metadata
    category: str = "general"
    requires_context: bool = True
    requires_privileged: bool = False
    skill_id: str | None = None  # If tool is from a skill

    def to_schema(self) -> LLMToolSpec:
        """Convert to LLMToolSpec Pydantic model (OpenAI function calling format)."""
        if isinstance(self.parameters, JSONSchema):
            params_dict = self.parameters.to_dict()
        elif isinstance(self.parameters, dict):
            params_dict = self.parameters
        else:
            params_dict = {"type": "object", "properties": {}}

        return LLMToolSpec(
            function=LLMFunctionSpec(
                name=self.name,
                description=self.description,
                parameters=LLMFunctionParameters.model_validate(params_dict),
            ),
        )


# =============================================================================
# TOOL REGISTRY
# =============================================================================


class ToolRegistry:
    """Centralized registry for all tools.

    Thread-safe tool storage with lookup and execution capabilities.

    Features:
    - Register/unregister tools
    - Lookup by name or pattern
    - Execute tools with context
    - Get schemas as LLMToolSpec models for LLM function calling
    """

    def __init__(self) -> None:
        self._tools: dict[str, ToolDefinition] = {}
        self._lock = threading.RLock()

    def register(
        self,
        name: str,
        handler: ToolHandler,
        description: str = "",
        parameters: JSONSchema | JsonObject | None = None,
        *,
        category: str = "general",
        requires_privileged: bool = False,
        skill_id: str | None = None,
    ) -> None:
        """Register a tool.

        Args:
            name: Tool name (e.g., "web.search")
            handler: Async function to execute
            description: Tool description
            parameters: JSONSchema model or dict for parameters
            category: Tool category
            requires_privileged: Whether tool requires elevated permissions
            skill_id: If tool is bundled with a skill
        """
        # Normalize parameters to JSONSchema model
        if parameters is None:
            params = JSONSchema(properties={})
        elif isinstance(parameters, dict):
            params = JSONSchema.from_dict(parameters)
        else:
            params = parameters

        tool = ToolDefinition(
            name=name,
            description=description,
            parameters=params,
            handler=handler,
            category=category,
            requires_privileged=requires_privileged,
            skill_id=skill_id,
        )

        with self._lock:
            if name in self._tools:
                logger.debug(f"Overwriting existing tool: {name}")
            self._tools[name] = tool
            logger.debug(f"Registered tool: {name}")

    def register_definition(self, tool: ToolDefinition) -> None:
        """Register a tool from a ToolDefinition."""
        with self._lock:
            if tool.name in self._tools:
                logger.debug(f"Overwriting existing tool: {tool.name}")
            self._tools[tool.name] = tool
            logger.debug(f"Registered tool: {tool.name}")

    def unregister(self, name: str) -> bool:
        """Unregister a tool.

        Returns
        -------
            True if tool was removed, False if not found
        """
        with self._lock:
            if name in self._tools:
                del self._tools[name]
                logger.debug(f"Unregistered tool: {name}")
                return True
            return False

    def get(self, name: str) -> ToolDefinition | None:
        """Get a tool by exact name."""
        return self._tools.get(name)

    def get_handler(self, name: str) -> ToolHandler | None:
        """Get a tool's handler by name."""
        tool = self._tools.get(name)
        return tool.handler if tool else None

    def list(self, pattern: str | None = None) -> list[ToolDefinition]:
        """List tools, optionally filtered by pattern.

        Args:
            pattern: Optional glob pattern (e.g., "web.*")

        Returns
        -------
            List of matching tools
        """
        if pattern is None:
            return list(self._tools.values())

        return [tool for name, tool in self._tools.items() if fnmatch.fnmatch(name, pattern)]

    def list_names(self) -> "list[str]":
        """List all tool names."""
        return list(self._tools.keys())

    def names(self) -> "list[str]":
        """List all tool names (alias for list_names)."""
        return self.list_names()

    def resolve_patterns(
        self,
        patterns: "list[str]",
        *,
        include_admin: bool = False,
    ) -> "list[ToolDefinition]":
        """Resolve tool patterns to actual tools.

        Args:
            patterns: List of patterns (e.g., ["web.*", "workspace.read_file"])
            include_admin: Include admin-category tools (default: hidden)

        Returns
        -------
            Deduplicated list of matching tools
        """
        resolved: dict[str, ToolDefinition] = {}

        for pattern in patterns:
            if "*" in pattern:
                # Glob pattern
                for name, tool in self._tools.items():
                    if fnmatch.fnmatch(name, pattern):
                        if tool.category == "admin" and not include_admin:
                            continue
                        resolved[name] = tool
            else:
                # Exact name
                tool = self._tools.get(pattern)
                if tool:
                    if tool.category == "admin" and not include_admin:
                        continue
                    resolved[pattern] = tool

        return list(resolved.values())

    def get_schemas(
        self,
        patterns: "list[str]",
        *,
        privileged: bool = False,
        include_admin: bool = False,
    ) -> "list[LLMToolSpec]":
        """Get LLMToolSpec models for tools matching patterns.

        Args:
            patterns: Tool patterns to match
            privileged: Include privileged tools
            include_admin: Include admin-category tools (default: hidden)

        Returns
        -------
            List of LLMToolSpec Pydantic models
        """
        tools = self.resolve_patterns(patterns, include_admin=include_admin)
        schemas: list[LLMToolSpec] = []

        for tool in tools:
            if tool.requires_privileged and not privileged:
                continue
            schemas.append(tool.to_schema())

        return schemas

    async def call(
        self,
        name: str,
        args: ToolArgs,
        context: object | None = None,
    ) -> ToolOutput:
        """Execute a tool by name.

        Args:
            name: Tool name
            args: ToolArgs Pydantic model
            context: Execution context

        Returns
        -------
            ToolOutput Pydantic model
        """
        tool = self._tools.get(name)
        if not tool:
            return ToolOutput(success=False, error=f"Tool not found: {name}")

        if not tool.handler:
            return ToolOutput(success=False, error=f"Tool has no handler: {name}")

        try:
            result = tool.handler(args, context)
            if hasattr(result, "__await__"):
                result = await result

            # Handlers MUST return ToolOutput. Dict fallback kept for edge cases.
            if isinstance(result, ToolOutput):
                return result
            if isinstance(result, dict):
                logger.warning(f"Tool {name} returned dict instead of ToolOutput")
                return ToolOutput.model_validate(result)
            return ToolOutput(success=True, data=str(result))
        except Exception as e:
            logger.exception(f"Tool call failed: {name}")
            return ToolOutput(success=False, error=str(e))

    def categories(self) -> "dict[str, list[str]]":
        """Get tools organized by category."""
        by_category: dict[str, list[str]] = {}
        for name, tool in self._tools.items():
            cat = tool.category or name.split(".")[0]
            by_category.setdefault(cat, []).append(name)
        return by_category

    def as_handler_dict(self) -> dict[str, ToolHandler]:
        """Return a ``{name: handler}`` dict for all tools with handlers.

        This is the bridge between ToolRegistry and Context.tools.
        Context expects ``dict[str, ToolHandler]``; this method produces it.
        """
        return {
            name: tool.handler for name, tool in self._tools.items() if tool.handler is not None
        }

    def clear(self) -> None:
        """Clear all registered tools."""
        with self._lock:
            self._tools.clear()

    def __len__(self) -> int:
        return len(self._tools)

    def __contains__(self, name: str) -> bool:
        return name in self._tools


# =============================================================================
# GLOBAL REGISTRY
# =============================================================================

_global_registry: ToolRegistry | None = None
_global_lock = threading.Lock()


def get_tool_registry() -> ToolRegistry:
    """Get the global tool registry.

    Typically set at bootstrap.  If called before initialization
    initialisation (e.g. in tests), a default instance is lazily created.
    """
    global _global_registry
    if _global_registry is None:
        with _global_lock:
            if _global_registry is None:
                _global_registry = ToolRegistry()
    return _global_registry


def set_tool_registry(registry: ToolRegistry) -> None:
    """Install *registry* as the global tool registry.

    Called at bootstrap to ensure the caller owns the instance.
    """
    global _global_registry
    _global_registry = registry


# =============================================================================
# EXPORTS
# =============================================================================

__all__ = [
    "ToolDefinition",
    "ToolHandler",
    "ToolRegistry",
    "get_tool_registry",
    "set_tool_registry",
]
