Bat/src/main/java/cc/fascinated/bat/service/CommandService.java

package cc.fascinated.bat.service;

import cc.fascinated.bat.Consts;
import cc.fascinated.bat.command.BatCommand;
import cc.fascinated.bat.command.Category;
import cc.fascinated.bat.common.EmbedUtils;
import cc.fascinated.bat.config.Config;
import cc.fascinated.bat.features.FeatureProfile;
import cc.fascinated.bat.model.BatGuild;
import cc.fascinated.bat.model.BatUser;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.log4j.Log4j2;
import net.dv8tion.jda.api.JDA;
import net.dv8tion.jda.api.Permission;
import net.dv8tion.jda.api.entities.Guild;
import net.dv8tion.jda.api.events.interaction.command.SlashCommandInteractionEvent;
import net.dv8tion.jda.api.hooks.ListenerAdapter;
import net.dv8tion.jda.api.interactions.commands.Command;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.DependsOn;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author Fascinated (fascinated7)
 */
@Service
@Log4j2(topic = "Command Service")
@Getter
@DependsOn("discordService")
public class CommandService extends ListenerAdapter {
    public static CommandService INSTANCE;

    /**
     * The registered commands
     */
    private final Map<String, BatCommand> commands = new HashMap<>();

    /**
     * The guild service to use
     */
    private final GuildService guildService;

    /**
     * The user service to use
     */
    private final UserService userService;

    @Autowired
    public CommandService(@NonNull GuildService guildService, @NonNull UserService userService) {
        INSTANCE = this;
        this.guildService = guildService;
        this.userService = userService;
        DiscordService.JDA.addEventListener(this);
    }

    /**
     * Registers a command
     *
     * @param command The command to register
     */
    public void registerCommand(@NonNull BatCommand command) {
        String commandName = command.getInfo().getName().toLowerCase();
        if (commands.get(commandName) != null) {
            return;
        }
        log.info("Registered command \"{}\"", command.getInfo().getName());
        commands.put(commandName, command);
    }

    /**
     * Registers all slash commands
     */
    public void registerSlashCommands() {
        log.info("Registering all slash commands");
        JDA jda = DiscordService.JDA;
        long before = System.currentTimeMillis();

        Guild adminGuild = jda.getGuildById(Config.INSTANCE.getAdminGuild());
        if (!Config.isProduction()) {
            if (adminGuild == null) {
                log.error("Unable to find the admin guild to register commands");
                return;
            }
            jda.retrieveCommands().complete().forEach(command -> jda.deleteCommandById(command.getId()).complete());
            List<Command> registeredCommands = adminGuild.updateCommands().addCommands(commands.values().stream().map(BatCommand::getCommandData).toList()).complete();
            log.info("Registered {} slash commands in {}ms (DEV MODE)", registeredCommands.size(), System.currentTimeMillis() - before);
            return;
        }

        // Unregister all commands that Discord has but we don't
        jda.retrieveCommands().complete().forEach(command -> {
            if (commands.containsKey(command.getName()) && commands.get(command.getName()).getInfo().isBotOwnerOnly()) {
                jda.deleteCommandById(command.getId()).complete(); // Unregister the command on Discord
                log.info("Unregistered hidden command \"{}\" from Discord", command.getName());
                return;
            }
            if (commands.containsKey(command.getName())) {
                return;
            }

            jda.deleteCommandById(command.getId()).complete(); // Unregister the command on Discord
            log.info("Unregistered unknown command \"{}\" from Discord", command.getName());
        });

        // Register all commands
        List<Command> discordCommands = jda.updateCommands().addCommands(commands.values().stream().filter(command -> !command.getInfo().isBotOwnerOnly())
                .map(BatCommand::getCommandData).toList()).complete();
        for (Command discordCommand : discordCommands) {
            commands.get(discordCommand.getName()).setSnowflake(discordCommand.getIdLong());
            if (!discordCommand.getSubcommands().isEmpty()) {
                for (Command.Subcommand subCommand : discordCommand.getSubcommands()) {
                    commands.get(discordCommand.getName()).getSubCommands().get(subCommand.getName()).setSnowflake(subCommand.getIdLong());
                }
            }
        }
        if (adminGuild != null) {
            adminGuild.updateCommands().addCommands(commands.values().stream().filter(command -> command.getInfo().isBotOwnerOnly())
                    .map(BatCommand::getCommandData).toList()).complete();
        } else {
            log.error("Unable to find the admin guild to register hidden commands");
        }
        log.info("Registered {} slash commands in {}ms", discordCommands.size(), System.currentTimeMillis() - before);
    }

    /**
     * Gets commands that are in a specific category
     *
     * @param category The category
     * @param includeGuildCommands If guild commands should be included
     * @return The commands
     */
    public List<BatCommand> getCommandsByCategory(Category category, boolean includeGuildCommands) {
        List<BatCommand> commands = new ArrayList<>();
        for (BatCommand command : this.commands.values()) {
            if (command.getInfo().getCategory() == category && (includeGuildCommands || command.getInfo().isUserInstall())) {
                commands.add(command);
            }
        }
        return commands;
    }

    /**
     * Gets a command by its class
     *
     * @param clazz The class of the command
     * @return The command
     */
    public BatCommand getCommand(Class<? extends BatCommand> clazz) {
        for (Map.Entry<String, BatCommand> entry : this.commands.entrySet()) {
            if (entry.getValue().getClass().equals(clazz)) {
                return entry.getValue();
            }
        }
        return null;
    }

    @Override
    public void onSlashCommandInteraction(@NotNull SlashCommandInteractionEvent event) {
        if (event.getUser().isBot()) {
            return;
        }
        long before = System.currentTimeMillis();
        BatCommand command = commands.get(event.getName());
        if (command == null) {
            return;
        }
        String subcommandName = event.getSubcommandName();
        if (subcommandName != null) { // Use the sub command if given
            command = command.getSubCommands().get(subcommandName);
        }
        BatUser user = userService.getUser(event.getUser().getId());
        BatGuild guild = event.getGuild() == null ? null : guildService.getGuild(event.getGuild().getId());

        // Only the bot owner can run this command
        if (command.getInfo().isBotOwnerOnly() && !user.getId().equalsIgnoreCase(Consts.BOT_OWNER)) {
            event.replyEmbeds(EmbedUtils.errorEmbed()
                    .setDescription("You do not have permission to execute this command")
                    .build()).setEphemeral(true).queue();
            return;
        }

        // Check if the command is guild only and if it was not ran inside a guild
        if (command.getInfo().isGuildOnly() && guild == null) {
            event.replyEmbeds(EmbedUtils.errorEmbed()
                    .setDescription("This command can only be executed in a guild")
                    .build()).setEphemeral(true).queue();
            return;
        }

        // Check if the feature is disabled in the guild
        if (guild != null) {
            FeatureProfile featureProfile = guild.getFeatureProfile();
            if (featureProfile.isFeatureDisabled(command.getFeature())) {
                event.replyEmbeds(EmbedUtils.errorEmbed()
                        .setDescription("The feature `%s` is disabled in this guild".formatted(command.getFeature().getName()))
                        .build()).setEphemeral(true).queue();
                return;
            }
        }

        // Check if the user has the required permissions
        if (guild != null && event.getMember() != null) {
            List<Permission> missingPermissions = new ArrayList<>();
            for (Permission permission : command.getInfo().getPermissions()) {
                if (!event.getMember().hasPermission(permission)) {
                    missingPermissions.add(permission);
                }
            }
            if (!missingPermissions.isEmpty()) {
                event.replyEmbeds(EmbedUtils.errorEmbed()
                                .setDescription("You are missing the following permissions to execute this command:\n" +
                                        missingPermissions.stream().map(perm -> "`" + perm.name() + "`").collect(Collectors.joining(", ")))
                                .build())
                        .setEphemeral(true)
                        .queue();
                return;
            }
        }

        // Execute the command
        try {
            command.execute(guild, user, event.getChannel(), event.getMember(), event.getInteraction());
            log.info("Executed command \"{}\" for user \"{}\" (took: {}ms)", command.getInfo().getName(), user.getName(), System.currentTimeMillis() - before);
        } catch (Exception ex) {
            log.error("An error occurred while executing command \"{}\"", command.getInfo().getName(), ex);
            event.replyEmbeds(EmbedUtils.genericInteractionError(ex).build()).setEphemeral(true).queue();
        }
    }
}