PluralKit/PluralKit.Bot/Interactive/BaseInteractive.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

using Autofac;

using Myriad.Rest.Types;
using Myriad.Rest.Types.Requests;
using Myriad.Types;

using NodaTime;

namespace PluralKit.Bot.Interactive
{
    public abstract class BaseInteractive
    {
        protected readonly Context _ctx;
        protected readonly List<Button> _buttons = new();
        protected readonly TaskCompletionSource _tcs = new();
        protected Message _message { get; private set; }
        protected bool _running;
        
        protected BaseInteractive(Context ctx)
        {
            _ctx = ctx;
        }

        public Duration Timeout { get; set; } = Duration.FromMinutes(5);

        protected Button AddButton(Func<InteractionContext, Task> handler, string? label = null, ButtonStyle style = ButtonStyle.Secondary, bool disabled = false)
        {
            var dispatch = _ctx.Services.Resolve<InteractionDispatchService>();
            var customId = dispatch.Register(handler, Timeout);
            
            var button = new Button
            {
                Label = label,
                Style = style,
                Disabled = disabled,
                CustomId = customId
            };
            _buttons.Add(button);
            return button;
        }

        protected async Task Update(InteractionContext ctx)
        {
            await ctx.Respond(InteractionResponse.ResponseType.UpdateMessage,
                new InteractionApplicationCommandCallbackData
                {
                    Components = GetComponents()
                });
        }

        protected async Task Finish(InteractionContext? ctx = null)
        {
            foreach (var button in _buttons) 
                button.Disabled = true;

            if (ctx != null)
                await Update(ctx);
            else 
                await _ctx.Rest.EditMessage(_message.ChannelId, _message.Id, new MessageEditRequest {
                    Components = GetComponents()
                });
            
            _tcs.TrySetResult();
        }

        protected async Task Send(string? content = null, Embed? embed = null, AllowedMentions? mentions = null)
        {
            _message = await _ctx.Rest.CreateMessage(_ctx.Channel.Id, new MessageRequest
            {
                Content = content,
                Embed = embed,
                AllowedMentions = mentions,
                Components = GetComponents()
            });
        }

        public MessageComponent[] GetComponents()
        {
            return new MessageComponent[]
            {
                new()
                {
                    Type = ComponentType.ActionRow,
                    Components = _buttons.Select(b => b.ToMessageComponent()).ToArray()
                }
            };
        }

        public void Setup(Context ctx)
        {
            var dispatch = ctx.Services.Resolve<InteractionDispatchService>();
            foreach (var button in _buttons) 
                button.CustomId = dispatch.Register(button.Handler, Timeout);
        }

        public abstract Task Start();
        
        public async Task Run()
        {
            if (_running)
                throw new InvalidOperationException("Action is already running");
            _running = true;

            await Start();
            
            var cts = new CancellationTokenSource(Timeout.ToTimeSpan());
            cts.Token.Register(() => _tcs.TrySetException(new TimeoutException("Action timed out")));

            try
            {
                await _tcs.Task;
            }
            finally
            {
                Cleanup();
            }
        }

        protected void Cleanup()
        {
            var dispatch = _ctx.Services.Resolve<InteractionDispatchService>();
            foreach (var button in _buttons) 
                dispatch.Unregister(button.CustomId!);
        }
    }
}