minecraft-kajetans-mod/src/main/java/me/km/permissions/ModCommandManager.java

package me.km.permissions;

import com.google.common.collect.Maps;
import com.mojang.brigadier.builder.ArgumentBuilder;
import com.mojang.brigadier.builder.RequiredArgumentBuilder;
import com.mojang.brigadier.tree.CommandNode;
import com.mojang.brigadier.tree.RootCommandNode;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import me.hammerle.snuviscript.code.ISnuviScheduler;
import me.km.events.CommandEvent;
import me.km.snuviscript.ScriptEvents;
import me.km.snuviscript.Scripts;
import net.minecraft.command.CommandSource;
import net.minecraft.command.Commands;
import net.minecraft.command.ICommandSource;
import net.minecraft.command.ISuggestionProvider;
import net.minecraft.command.arguments.SuggestionProviders;
import net.minecraft.entity.Entity;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraft.entity.player.ServerPlayerEntity;
import net.minecraft.network.play.server.SCommandListPacket;

public class ModCommandManager extends Commands
{
    private final HashSet<String> otherCommands = new HashSet<>();
    
    private final HashMap<String, Command> commands = new HashMap<>();
    private final PermissionManager perms;
    private final ScriptEvents events;
    private final Scripts scripts;
    
    private final ArrayList<CommandNode> customNodes = new ArrayList<>();

    public ModCommandManager(boolean isDedicatedServer, PermissionManager perms, ScriptEvents events, Scripts scripts, ISnuviScheduler scheduler)
    {
        super(isDedicatedServer);
        getDispatcher().getRoot().getChildren().forEach(c -> 
        {
            otherCommands.add(c.getName());
        });
        scheduler.scheduleTask(() -> 
        {
            getDispatcher().getRoot().getChildren().forEach(c -> 
            {
                if(otherCommands.add(c.getName()))
                {
                    perms.addOtherGroupPermission(c.getName());
                }
            }); 
        });
        this.perms = perms;
        this.events = events;
        this.scripts = scripts;
    }
    
    public void clearCustomNodes()
    {
        customNodes.clear();
    }
    
    public void addCustomNodes(CommandNode node)
    {
        customNodes.add(node);
    }
    
    public void registerCommand(Command command)
    {
        commands.put(command.getName(), command);
        for(String alias : command.getAliases())
        {
            commands.put(alias, command);
        }
    }

    private String getCommandName(String rawCommand)
    {
        if(rawCommand.isEmpty())
        {
            return "";
        }
        int index = rawCommand.indexOf(' ');
        return rawCommand.substring(rawCommand.charAt(0) == '/' ? 1 : 0, index == -1 ? rawCommand.length() : index);
    }
    
    private static String[] getArguments(String rawCommand)
    {
        int old = rawCommand.indexOf(' ') + 1;
        if(old == 0)
        {
            return new String[0];
        }
        int pos = old;
        ArrayList<String> list = new ArrayList<>();
        while(pos < rawCommand.length())
        {
            char c = rawCommand.charAt(pos);
            switch(c)
            {
                case ' ':
                    if(pos - old > 0)
                    {
                        list.add(rawCommand.substring(old, pos));
                    }
                    old = pos + 1;
                    break;
                case '"':
                    if(pos - old > 0)
                    {
                        list.add(rawCommand.substring(old, pos));
                    }
                    old = pos + 1;
                    pos = old;
                    while(pos < rawCommand.length() && rawCommand.charAt(pos) != '"')
                    {
                        pos++;
                    }
                    list.add(rawCommand.substring(old, pos));
                    old = pos + 1;
                    break;
            }
            pos++;
        }
        if(pos - old > 0)
        {
            list.add(rawCommand.substring(old, pos));
        }
        return list.toArray(new String[list.size()]);
    }

    private ICommandSource getSource(CommandSource cs)
    {
        if(cs.getEntity() != null)
        {
            return cs.getEntity();
        }
        return cs.getServer();
    }
    
    private String lowerRawCommand(String raw)
    {
        int index = raw.indexOf(" ");
        if(index == -1)
        {
            return raw.toLowerCase();
        }
        return raw.substring(0, index).toLowerCase() + raw.substring(index);
    }
    
    @Override
    public int handleCommand(CommandSource cs, String rawCommand)
    { 
        rawCommand = lowerRawCommand(rawCommand);
        String commandName = getCommandName(rawCommand);
        
        Command command = commands.get(commandName);
        if(command != null)
        {
            if(perms.hasPermission(cs, command.getName()))
            {
                command.execute(getSource(cs), getArguments(rawCommand));
                return 1;
            }
            events.onMissingPermission(getSource(cs), command.getName());
            return 0;
        }
        
        if(scripts.isRegisteredScriptCommand(commandName))
        {
            Entity ent = cs.getEntity();
            if(ent != null && ent instanceof PlayerEntity)
            {
                events.onCustomCommand((PlayerEntity) ent, commandName, getArguments(rawCommand));
                return 1;
            }
            events.onCustomCommand(null, commandName, getArguments(rawCommand));
            return 0;
        }
        
        if(otherCommands.contains(commandName))
        {
            if(perms.hasPermission(cs, commandName))
            {
                Entity ent = cs.getEntity();
                if(ent != null && ent instanceof PlayerEntity)
                {
                    CommandEvent e = new CommandEvent((PlayerEntity) ent, commandName);
                    events.onCommand(e);
                    if(e.isCanceled())
                    {
                        return 0;
                    }
                }
                return super.handleCommand(cs, rawCommand);
            }
            events.onMissingPermission(getSource(cs), commandName);
            return 0;
        }
        
        events.onMissingCommand(getSource(cs), commandName);
        return 0;
    }

    @Override
    public void send(ServerPlayerEntity player)
    {
        Map<CommandNode<CommandSource>, CommandNode<ISuggestionProvider>> map = Maps.newHashMap();
        RootCommandNode<ISuggestionProvider> rootcommandnode = new RootCommandNode<>();
        map.put(getDispatcher().getRoot(), rootcommandnode);
        this.commandSourceNodesToSuggestionNodes(true, getDispatcher().getRoot(), rootcommandnode, player.getCommandSource(), map);
        
        for(CommandNode node : customNodes)
        {
            rootcommandnode.addChild(node);
        }
        
        player.connection.sendPacket(new SCommandListPacket(rootcommandnode));
    }

    private void commandSourceNodesToSuggestionNodes(boolean first, 
            CommandNode<CommandSource> node, 
            CommandNode<ISuggestionProvider> suggestion, 
            CommandSource source, 
            Map<CommandNode<CommandSource>, CommandNode<ISuggestionProvider>> map)
    {
        for(CommandNode<CommandSource> commandnode : node.getChildren())
        {
            if((first && perms.hasPermission(source, commandnode.getName()) || (!first && commandnode.canUse(source))))
            {
                ArgumentBuilder<ISuggestionProvider, ?> argumentbuilder = (ArgumentBuilder) commandnode.createBuilder();
                argumentbuilder.requires(a -> true);
                if(argumentbuilder.getCommand() != null)
                {
                    argumentbuilder.executes(a -> 0);
                }

                if(argumentbuilder instanceof RequiredArgumentBuilder)
                {
                    RequiredArgumentBuilder<ISuggestionProvider, ?> requiredargumentbuilder = (RequiredArgumentBuilder) argumentbuilder;
                    if(requiredargumentbuilder.getSuggestionsProvider() != null)
                    {
                        requiredargumentbuilder.suggests(SuggestionProviders.ensureKnown(requiredargumentbuilder.getSuggestionsProvider()));
                    }
                }

                if(argumentbuilder.getRedirect() != null)
                {
                    argumentbuilder.redirect(map.get(argumentbuilder.getRedirect()));
                }

                CommandNode<ISuggestionProvider> commandNode = argumentbuilder.build();
                map.put(commandnode, commandNode);
                suggestion.addChild(commandNode);
                if(!commandnode.getChildren().isEmpty())
                {
                    this.commandSourceNodesToSuggestionNodes(false, commandnode, commandNode, source, map);
                }
            }
        }
    }
}