PluralKit/PluralKit.Bot/CommandSystem/ContextEntityArgumentsExt.cs

using System;
using System.Threading.Tasks;

using Myriad.Extensions;
using Myriad.Types;

using PluralKit.Bot.Utils;
using PluralKit.Core;

namespace PluralKit.Bot
{
    public static class ContextEntityArgumentsExt
    {
        public static async Task<User> MatchUser(this Context ctx)
        {
            var text = ctx.PeekArgument();
            if (text.TryParseMention(out var id))
                return await ctx.Cache.GetOrFetchUser(ctx.Rest, id);

            return null;
        }

        public static bool MatchUserRaw(this Context ctx, out ulong id)
        {
            id = 0;

            var text = ctx.PeekArgument();
            if (text.TryParseMention(out var mentionId))
                id = mentionId;

            return id != 0;
        }

        public static Task<PKSystem> PeekSystem(this Context ctx) => ctx.MatchSystemInner();

        public static async Task<PKSystem> MatchSystem(this Context ctx)
        {
            var system = await ctx.MatchSystemInner();
            if (system != null) ctx.PopArgument();
            return system;
        }

        private static async Task<PKSystem> MatchSystemInner(this Context ctx)
        {
            var input = ctx.PeekArgument();

            // System references can take three forms:
            // - The direct user ID of an account connected to the system
            // - A @mention of an account connected to the system (<@uid>)
            // - A system hid

            // Direct IDs and mentions are both handled by the below method:
            if (input.TryParseMention(out var id))
                return await ctx.Repository.GetSystemByAccount(id);

            // Finally, try HID parsing
            var system = await ctx.Repository.GetSystemByHid(input);
            return system;
        }

        public static async Task<PKMember> PeekMember(this Context ctx, SystemId? restrictToSystem = null)
        {
            var input = ctx.PeekArgument();

            // Member references can have one of three forms, depending on
            // whether you're in a system or not:
            // - A member hid
            // - A textual name of a member *in your own system*
            // - a textual display name of a member *in your own system*

            // Skip name / display name matching if the user does not have a system
            // or if they specifically request by-HID matching
            if (ctx.System != null && !ctx.MatchFlag("id", "by-id"))
            {
                // First, try finding by member name in system
                if (await ctx.Repository.GetMemberByName(ctx.System.Id, input) is PKMember memberByName)
                    return memberByName;

                // And if that fails, we try finding a member with a display name matching the argument from the system
                if (ctx.System != null && await ctx.Repository.GetMemberByDisplayName(ctx.System.Id, input) is PKMember memberByDisplayName)
                    return memberByDisplayName;
            }

            // Finally (or if by-HID lookup is specified), try member HID parsing:
            if (await ctx.Repository.GetMemberByHid(input, restrictToSystem) is PKMember memberByHid)
                return memberByHid;

            // We didn't find anything, so we return null.
            return null;
        }

        /// <summary>
        /// Attempts to pop a member descriptor from the stack, returning it if present. If a member could not be
        /// resolved by the next word in the argument stack, does *not* touch the stack, and returns null.
        /// </summary>
        public static async Task<PKMember> MatchMember(this Context ctx, SystemId? restrictToSystem = null)
        {
            // First, peek a member
            var member = await ctx.PeekMember(restrictToSystem);

            // If the peek was successful, we've used up the next argument, so we pop that just to get rid of it.
            if (member != null) ctx.PopArgument();

            // Finally, we return the member value.
            return member;
        }

        public static async Task<PKGroup> PeekGroup(this Context ctx, SystemId? restrictToSystem = null)
        {
            var input = ctx.PeekArgument();

            // see PeekMember for an explanation of the logic used here

            if (ctx.System != null && !ctx.MatchFlag("id", "by-id"))
            {
                if (await ctx.Repository.GetGroupByName(ctx.System.Id, input) is { } byName)
                    return byName;
                if (await ctx.Repository.GetGroupByDisplayName(ctx.System.Id, input) is { } byDisplayName)
                    return byDisplayName;
            }

            if (await ctx.Repository.GetGroupByHid(input, restrictToSystem) is { } byHid)
                return byHid;

            return null;
        }

        public static async Task<PKGroup> MatchGroup(this Context ctx, SystemId? restrictToSystem = null)
        {
            var group = await ctx.PeekGroup(restrictToSystem);
            if (group != null) ctx.PopArgument();
            return group;
        }

        public static string CreateNotFoundError(this Context ctx, string entity, string input)
        {
            var isIDOnlyQuery = ctx.System == null || ctx.MatchFlag("id", "by-id");

            if (isIDOnlyQuery)
            {
                if (input.Length == 5)
                    return $"{entity} with ID \"{input}\" not found.";
                else
                    return $"{entity} not found. Note that a {entity.ToLower()} ID is 5 characters long.";
            }
            else
            {
                if (input.Length == 5)
                    return $"{entity} with ID or name \"{input}\" not found.";
                else
                    return $"{entity} with name \"{input}\" not found. Note that a {entity.ToLower()} ID is 5 characters long.";
            }
        }

        public static Task<Channel> MatchChannel(this Context ctx)
        {
            if (!MentionUtils.TryParseChannel(ctx.PeekArgument(), out var id))
                return Task.FromResult<Channel>(null);

            if (!ctx.Cache.TryGetChannel(id, out var channel))
                return Task.FromResult<Channel>(null);

            if (!DiscordUtils.IsValidGuildChannel(channel))
                return Task.FromResult<Channel>(null);

            ctx.PopArgument();
            return Task.FromResult(channel);
        }

        public static Guild MatchGuild(this Context ctx)
        {
            try
            {
                var id = ulong.Parse(ctx.PeekArgument());
                ctx.Cache.TryGetGuild(id, out var guild);
                if (guild != null)
                    ctx.PopArgument();

                return guild;
            }
            catch (FormatException)
            {
                return null;
            }
        }
    }
}