    001package co.aikar.commands;
    002
    003import co.aikar.commands.annotation.Author;
    004import co.aikar.commands.annotation.CrossGuild;
    005import co.aikar.commands.annotation.Optional;
    006import co.aikar.commands.annotation.SelfUser;
    007import net.dv8tion.jda.core.JDA;
    008import net.dv8tion.jda.core.entities.ChannelType;
    009import net.dv8tion.jda.core.entities.Guild;
    010import net.dv8tion.jda.core.entities.Message;
    011import net.dv8tion.jda.core.entities.MessageChannel;
    012import net.dv8tion.jda.core.entities.Role;
    013import net.dv8tion.jda.core.entities.TextChannel;
    014import net.dv8tion.jda.core.entities.User;
    015import net.dv8tion.jda.core.events.message.MessageReceivedEvent;
    016
    017import java.util.List;
    018
    019// TODO: Message Keys !!!
    020public class JDACommandContexts extends CommandContexts<JDACommandExecutionContext> {
    021    private final JDACommandManager manager;
    022    private final JDA jda;
    023
    024    public JDACommandContexts(JDACommandManager manager) {
    025        super(manager);
    026        this.manager = manager;
    027        this.jda = this.manager.getJDA();
    028        this.registerIssuerOnlyContext(JDACommandEvent.class, CommandExecutionContext::getIssuer);
    029        this.registerIssuerOnlyContext(MessageReceivedEvent.class, c -> c.getIssuer().getIssuer());
    030        this.registerIssuerOnlyContext(Message.class, c -> c.issuer.getIssuer().getMessage());
    031        this.registerIssuerOnlyContext(ChannelType.class, c -> c.issuer.getIssuer().getChannelType());
    032        this.registerIssuerOnlyContext(JDA.class, c -> jda);
    033        this.registerIssuerOnlyContext(Guild.class, c -> {
    034            MessageReceivedEvent event = c.getIssuer().getIssuer();
    035            if (event.isFromType(ChannelType.PRIVATE) && !c.isOptional()) {
    036                throw new InvalidCommandArgument("This command can only be executed in a Guild.", false);
    037            } else {
    038                return event.getGuild();
    039            }
    040        });
    041        this.registerIssuerAwareContext(MessageChannel.class, c -> {
    042            if (c.hasAnnotation(Author.class)) {
    043                return c.issuer.getIssuer().getChannel();
    044            }
    045            boolean isCrossGuild = c.hasAnnotation(CrossGuild.class);
    046            String argument = c.popFirstArg(); // we pop because we are only issuer aware if we are annotated
    047            MessageChannel channel = null;
    048            if (argument.startsWith("<#")) {
    049                String id = argument.substring(2, argument.length() - 1);
    050                channel = isCrossGuild ? jda.getTextChannelById(id) : c.issuer.getIssuer().getGuild().getTextChannelById(id);
    051            } else {
    052                List<TextChannel> channelList = isCrossGuild ? jda.getTextChannelsByName(argument, true) :
    053                        c.issuer.getEvent().getGuild().getTextChannelsByName(argument, true);
    054                if (channelList.size() > 1) {
    055                    throw new InvalidCommandArgument("Too many channels were found with the given name. Try with the `#channelname` syntax.", false);
    056                } else if (channelList.size() == 1) {
    057                    channel = channelList.get(0);
    058                }
    059            }
    060            if (channel == null) {
    061                throw new InvalidCommandArgument("Couldn't find a channel with that name or ID.");
    062            }
    063            return channel;
    064        });
    065        this.registerIssuerAwareContext(User.class, c -> {
    066            if (c.hasAnnotation(SelfUser.class)) {
    067                return jda.getSelfUser();
    068            }
    069            String arg = c.getFirstArg();
    070            if (c.isOptional() && (arg == null || arg.isEmpty())) {
    071                return null;
    072            }
    073            arg = c.popFirstArg(); // we pop because we are only issuer aware if we are annotated
    074            User user = null;
    075            if (arg.startsWith("<@!")) { // for some reason a ! is added when @'ing and clicking their name.
    076                user = jda.getUserById(arg.substring(3, arg.length() - 1));
    077            } else if (arg.startsWith("<@")) { // users can /also/ be mentioned like this...
    078                user = jda.getUserById(arg.substring(2, arg.length() - 1));
    079            } else {
    080                List<User> users = jda.getUsersByName(arg, true);
    081                if (users.size() > 1) {
    082                    throw new InvalidCommandArgument("Too many users were found with the given name. Try with the `@username#0000` syntax.", false);
    083                }
    084                if (!users.isEmpty()) {
    085                    user = users.get(0);
    086                }
    087            }
    088            if (user == null) {
    089                throw new InvalidCommandArgument("Could not find a user with that name or ID.");
    090            }
    091            return user;
    092        });
    093        this.registerContext(Role.class, c -> {
    094            boolean isCrossGuild = c.hasAnnotation(CrossGuild.class);
    095            String arg = c.popFirstArg();
    096            Role role = null;
    097            if (arg.startsWith("<@&")) {
    098                String id = arg.substring(3, arg.length() - 1);
    099                role = isCrossGuild ? jda.getRoleById(id) : c.issuer.getIssuer().getGuild().getRoleById(id);
    100            } else {
    101                List<Role> roles = isCrossGuild ? jda.getRolesByName(arg, true)
    102                        : c.issuer.getIssuer().getGuild().getRolesByName(arg, true);
    103                if (roles.size() > 1) {
    104                    throw new InvalidCommandArgument("Too many roles were found with the given name. Try with the `@role` syntax.", false);
    105                }
    106                if (!roles.isEmpty()) {
    107                    role = roles.get(0);
    108                }
    109            }
    110            if (role == null) {
    111                throw new InvalidCommandArgument("Could not find a role with that name or ID.");
    112            }
    113            return role;
    114        });
    115    }
    116}