discordbot/command_sync_utils.py

"""
Utility module for syncing commands with Discord.

This module provides functions for syncing both global and guild-specific commands,
allowing for more flexibility in command registration and supporting commands in
both guild and private contexts (DMs and private channels).
"""

import discord
from discord import app_commands
from discord.ext import commands
import logging
import command_customization
from typing import Dict, List, Optional, Union, Any
from command_context import AppCommandContext

# Configure logging
logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

async def sync_global_and_guild_commands(
    bot: commands.Bot,
    global_only: bool = False,
    guild_only: bool = False,
    private_command_names: List[str] = None
) -> Dict[str, Any]:
    """
    Sync commands with Discord, first to guilds and then globally.

    This function implements the following strategy:
    1. First, sync guild-specific commands to every guild (as before)
    2. Then, sync global commands with dm_permission=True and guild_only=False

    Args:
        bot: The bot instance
        global_only: If True, only sync global commands
        guild_only: If True, only sync guild-specific commands
        private_command_names: List of command names that should work in private contexts

    Returns:
        A dictionary containing the results of the sync operations
    """
    results = {
        "global": [],
        "guild": {}
    }

    try:
        # Get all commands from the command tree
        all_commands = bot.tree.get_commands()
        command_names = [cmd.name for cmd in all_commands]
        log.info(f"Found {len(all_commands)} commands in tree: {', '.join(command_names)}")

        # If private_command_names is not provided, use an empty list
        if private_command_names is None:
            private_command_names = []

        # First, sync guild-specific commands if requested
        if not global_only:
            log.info("Syncing guild-specific commands to all guilds...")
            # Use the command_customization module to handle guild-specific commands
            guild_syncs = await command_customization.register_all_guild_commands(bot)
            results["guild"] = guild_syncs

            total_guild_syncs = sum(len(cmds) for cmds in guild_syncs.values())
            log.info(f"Synced commands for {len(guild_syncs)} guilds with a total of {total_guild_syncs} customized commands")

        # Now, prepare global commands with appropriate context settings
        if not guild_only and private_command_names:
            log.info("Preparing global commands for private contexts...")

            # Clear the command tree for global commands
            bot.tree.clear_commands()

            # Add back only the commands that should work in private contexts
            from command_context import AppCommandContext, set_command_context
            for cmd_name in private_command_names:
                # Find the command in the original list
                cmd = next((c for c in all_commands if c.name == cmd_name), None)
                if cmd:
                    log.info(f"Setting command {cmd_name} to work in private contexts")
                    # Set the command to work in private contexts
                    set_command_context(cmd, AppCommandContext.ALL)
                    # Add the command back to the tree
                    bot.tree.add_command(cmd)
                else:
                    log.warning(f"Command {cmd_name} not found in command tree")

            # Sync the global commands
            log.info("Syncing global commands for private contexts...")
            global_synced = await bot.tree.sync()
            results["global"] = global_synced

            # List the synced commands
            global_command_names = [cmd.name for cmd in global_synced]
            log.info(f"Synced {len(global_synced)} global command(s) for private contexts: {', '.join(global_command_names)}")

            # Restore the original command tree
            bot.tree.clear_commands()
            for cmd in all_commands:
                bot.tree.add_command(cmd)

    except Exception as e:
        log.error(f"Failed to sync commands: {e}")
        import traceback
        traceback.print_exc()

    return results

def register_global_command(
    bot: commands.Bot,
    name: str,
    description: str,
    callback: Any,
    context: AppCommandContext = AppCommandContext.ALL,
    **kwargs
) -> app_commands.Command:
    """
    Register a global command with the specified context.

    Args:
        bot: The bot instance
        name: The name of the command
        description: The description of the command
        callback: The function to call when the command is used
        context: The context to set for the command
        **kwargs: Additional arguments to pass to the command constructor

    Returns:
        The created command
    """
    from command_context import create_global_command

    return create_global_command(
        bot.tree,
        name=name,
        description=description,
        callback=callback,
        context=context,
        **kwargs
    )

def set_command_contexts(bot: commands.Bot, contexts: Dict[str, AppCommandContext]) -> None:
    """
    Set the contexts for multiple commands.

    Args:
        bot: The bot instance
        contexts: A dictionary mapping command names to contexts
    """
    from command_context import set_command_context

    for cmd in bot.tree.get_commands():
        if cmd.name in contexts:
            set_command_context(cmd, contexts[cmd.name])
            log.info(f"Set context for command '{cmd.name}' to {contexts[cmd.name]}")