"""
Tool Registry for Conditional Tool Registration

This module provides a registry system that allows tools to be conditionally registered
based on tier configuration, replacing direct @server.tool() decorators.
"""

import logging
from typing import Set, Optional, Callable

logger = logging.getLogger(__name__)

# Global registry of enabled tools
_enabled_tools: Optional[Set[str]] = None

def set_enabled_tools(tool_names: Optional[Set[str]]):
    """Set the globally enabled tools."""
    global _enabled_tools
    _enabled_tools = tool_names

def get_enabled_tools() -> Optional[Set[str]]:
    """Get the set of enabled tools, or None if all tools are enabled."""
    return _enabled_tools

def is_tool_enabled(tool_name: str) -> bool:
    """Check if a specific tool is enabled."""
    if _enabled_tools is None:
        return True  # All tools enabled by default
    return tool_name in _enabled_tools

def conditional_tool(server, tool_name: str):
    """
    Decorator that conditionally registers a tool based on the enabled tools set.
    
    Args:
        server: The FastMCP server instance
        tool_name: The name of the tool to register
    
    Returns:
        Either the registered tool decorator or a no-op decorator
    """
    def decorator(func: Callable) -> Callable:
        if is_tool_enabled(tool_name):
            logger.debug(f"Registering tool: {tool_name}")
            return server.tool()(func)
        else:
            logger.debug(f"Skipping tool registration: {tool_name}")
            return func
    
    return decorator

def wrap_server_tool_method(server):
    """
    Track tool registrations and filter them post-registration.
    """
    original_tool = server.tool
    server._tracked_tools = []
    
    def tracking_tool(*args, **kwargs):
        original_decorator = original_tool(*args, **kwargs)
        
        def wrapper_decorator(func: Callable) -> Callable:
            tool_name = func.__name__
            server._tracked_tools.append(tool_name)
            # Always apply the original decorator to register the tool
            return original_decorator(func)
        
        return wrapper_decorator
    
    server.tool = tracking_tool

def filter_server_tools(server):
    """Remove disabled tools from the server after registration."""
    enabled_tools = get_enabled_tools()
    if enabled_tools is None:
        return
    
    tools_removed = 0
    
    # Access FastMCP's tool registry via _tool_manager._tools
    if hasattr(server, '_tool_manager'):
        tool_manager = server._tool_manager
        if hasattr(tool_manager, '_tools'):
            tool_registry = tool_manager._tools
            
            tools_to_remove = []
            for tool_name in list(tool_registry.keys()):
                if not is_tool_enabled(tool_name):
                    tools_to_remove.append(tool_name)
            
            for tool_name in tools_to_remove:
                del tool_registry[tool_name]
                tools_removed += 1
    
    if tools_removed > 0:
        logger.info(f"🔧 Tool tier filtering: removed {tools_removed} tools, {len(enabled_tools)} enabled")