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}