From 795a69ca7d58dea786e30327f1d18f7c722e9b0f Mon Sep 17 00:00:00 2001
From: diamondburned <datutbrus@gmail.com>
Date: Thu, 14 May 2020 18:52:15 -0700
Subject: [PATCH] Bot: Added EditableCommands

---
 README.md                     | 10 --------
 _example/advanced_bot/main.go |  1 +
 bot/ctx.go                    |  5 ++++
 bot/ctx_call.go               | 48 +++++++++++++++++++++++++++++------
 bot/subcommand.go             |  1 +
 5 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 8210233..477f3c3 100644
--- a/README.md
+++ b/README.md
@@ -63,16 +63,6 @@ custom remote or local state storage.
 things in the state, which is useful for keeping it updated.
 - No code generation: just so the library is a lot easier to maintain.
 
-## You-should-knows
-
-- ~~The bot will fatally exit if it fails to reconnect to the Gateway after a
-certain amount of times. This is changeable in `gateway.WSFatal`, or
-`(*Gateway).FatalLog`.~~
-- ~~The bot will error out if the initial connection fails. However,
-reconnections will be retried forever until it succeeds.~~ This is no longer
-true. The bot will retry until `WSRetries` is reached, then the error will go 
-to `(*Gateway).FatalError` or `(*Gateway).Wait()`.
-
 ## Testing
 
 The package includes integration tests that require `$BOT_TOKEN`. To run these
diff --git a/_example/advanced_bot/main.go b/_example/advanced_bot/main.go
index d8e8ab4..bc5fc41 100644
--- a/_example/advanced_bot/main.go
+++ b/_example/advanced_bot/main.go
@@ -19,6 +19,7 @@ func main() {
 
 	wait, err := bot.Start(token, commands, func(ctx *bot.Context) error {
 		ctx.HasPrefix = bot.NewPrefix("!", "~")
+		ctx.EditableCommands = true
 
 		// Subcommand demo, but this can be in another package.
 		ctx.MustRegisterSubcommand(&Debug{})
diff --git a/bot/ctx.go b/bot/ctx.go
index cf4fdad..4b71d16 100644
--- a/bot/ctx.go
+++ b/bot/ctx.go
@@ -124,6 +124,11 @@ type Context struct {
 	// MessageCreate events.
 	ReplyError bool
 
+	// EditableCommands when true will also listen for MessageUpdateEvent and
+	// treat them as newly created messages. This is convenient if you want
+	// to quickly edit a message and re-execute the command.
+	EditableCommands bool
+
 	// Subcommands contains all the registered subcommands. This is not
 	// exported, as it shouldn't be used directly.
 	subcommands []*Subcommand
diff --git a/bot/ctx_call.go b/bot/ctx_call.go
index 2f58786..9415bc1 100644
--- a/bot/ctx_call.go
+++ b/bot/ctx_call.go
@@ -58,19 +58,51 @@ func (ctx *Context) callCmd(ev interface{}) (bottomError error) {
 		}
 	}
 
+	var msc *gateway.MessageCreateEvent
+
 	// We call the messages later, since we want MessageCreate middlewares to
 	// run as well.
-	if evT == typeMessageCreate {
-		// safe assertion always
-		err := ctx.callMessageCreate(ev.(*gateway.MessageCreateEvent), evV)
-		// There's no need for an errNoBreak here, as the method already checked
-		// for that.
-		if err != nil {
-			bottomError = err
+	switch {
+	case evT == typeMessageCreate:
+		msc = ev.(*gateway.MessageCreateEvent)
+
+	case evT == typeMessageUpdate && ctx.EditableCommands:
+		up := ev.(*gateway.MessageUpdateEvent)
+		// Message updates could have empty contents when only their embeds are
+		// filled. We don't need that here.
+		if up.Content == "" {
+			return nil
 		}
+
+		// Query the updated message.
+		m, err := ctx.Store.Message(up.ChannelID, up.ID)
+		if err != nil {
+			// It's probably safe to ignore this.
+			return nil
+		}
+
+		// Treat the message update as a message create event to avoid breaking
+		// changes.
+		msc = &gateway.MessageCreateEvent{Message: *m, Member: up.Member}
+
+		// Fill up member, if available.
+		if m.GuildID.Valid() && up.Member == nil {
+			if mem, err := ctx.Store.Member(m.GuildID, m.Author.ID); err == nil {
+				msc.Member = mem
+			}
+		}
+
+		// Update the reflect value as well.
+		evV = reflect.ValueOf(msc)
+
+	default:
+		// Unknown event, return.
+		return nil
 	}
 
-	return
+	// There's no need for an errNoBreak here, as the method already checked
+	// for that.
+	return ctx.callMessageCreate(msc, evV)
 }
 
 func (ctx *Context) callMessageCreate(mc *gateway.MessageCreateEvent, value reflect.Value) error {
diff --git a/bot/subcommand.go b/bot/subcommand.go
index e82d141..7b76b9c 100644
--- a/bot/subcommand.go
+++ b/bot/subcommand.go
@@ -12,6 +12,7 @@ import (
 
 var (
 	typeMessageCreate = reflect.TypeOf((*gateway.MessageCreateEvent)(nil))
+	typeMessageUpdate = reflect.TypeOf((*gateway.MessageUpdateEvent)(nil))
 
 	typeString = reflect.TypeOf("")
 	typeEmbed  = reflect.TypeOf((*discord.Embed)(nil))