edit.go

package strut

import (
	"time"

	"github.com/disgoorg/disgo/discord"
	"github.com/disgoorg/disgo/events"
	"github.com/disgoorg/snowflake/v2"
)

// defaultEditTTL is how long an invocation stays editable by default.
const defaultEditTTL = 10 * time.Minute

// EditTrackerConfig tunes how message edits are followed.
type EditTrackerConfig struct {
	// TTL is how long an invocation stays editable. Zero means ten minutes.
	TTL time.Duration
	// ExecuteUntrackedEdits runs a command when an edit turns a message into
	// a valid invocation, even though the original was not one.
	ExecuteUntrackedEdits bool
	// IgnoreEditsIfNotResponded skips edits to invocations the bot never
	// answered.
	IgnoreEditsIfNotResponded bool
	// Store keeps the invocation to reply mapping. Nil uses an in-process
	// store, which is usually right: a guild always lands on one shard.
	Store Store[snowflake.ID, TrackedReply]
}

// TrackedReply is the answer strut gave to one invocation.
type TrackedReply struct {
	// Response is nil when the command produced no reply.
	Response *discord.Message
	// DeleteWithInvocation records Meta.TrackDeletion, so deleting the
	// invocation removes the answer with it.
	DeleteWithInvocation bool
}

// EditTracker remembers which bot message answered which invocation, so an
// edited command can update its original reply instead of sending a new one.
type EditTracker struct {
	cfg   EditTrackerConfig
	store Store[snowflake.ID, TrackedReply]
}

// NewEditTracker returns a tracker with the given configuration.
func NewEditTracker(cfg EditTrackerConfig) *EditTracker {
	if cfg.TTL <= 0 {
		cfg.TTL = defaultEditTTL
	}
	store := cfg.Store
	if store == nil {
		store = NewMemoryStore[snowflake.ID, TrackedReply]()
	}
	return &EditTracker{cfg: cfg, store: store}
}

// track records how an invocation was answered.
func (t *EditTracker) track(invocation snowflake.ID, reply TrackedReply) {
	t.store.Put(invocation, reply, t.cfg.TTL)
}

// lookup returns the answer to an invocation, and whether it was tracked at
// all. A tracked invocation with no reply returns a zero TrackedReply.
func (t *EditTracker) lookup(invocation snowflake.ID) (TrackedReply, bool) {
	return t.store.Get(invocation)
}

// forget drops an invocation, used when its message is deleted.
func (t *EditTracker) forget(invocation snowflake.ID) (TrackedReply, bool) {
	return t.store.Remove(invocation)
}

// onMessageUpdate re-runs a command whose invocation was edited.
func (f *Framework[D]) onMessageUpdate(ev *events.MessageUpdate) {
	p := f.opts.Prefix
	if p == nil || p.EditTracker == nil {
		return
	}
	if ev.Message.Author.Bot && !p.AllowBots {
		return
	}

	tracked, ok := p.EditTracker.lookup(ev.MessageID)

	switch {
	case ok && tracked.Response == nil && p.EditTracker.cfg.IgnoreEditsIfNotResponded:
		return
	case !ok && !p.EditTracker.cfg.ExecuteUntrackedEdits:
		return
	}

	msg := &events.MessageCreate{GenericMessage: ev.GenericMessage}

	node, lex, found := f.resolvePrefix(msg)
	if !found {
		return
	}
	// Only commands that opted in re-run; the rest keep their first answer.
	if node.inv != nil && !node.inv.meta.InvokeOnEdit {
		return
	}
	f.dispatchPrefix(node, msg, lex, tracked.Response)
}

// onMessageDelete removes the bot's answer when the invocation is deleted.
func (f *Framework[D]) onMessageDelete(ev *events.MessageDelete) {
	p := f.opts.Prefix
	if p == nil || p.EditTracker == nil {
		return
	}

	tracked, ok := p.EditTracker.forget(ev.MessageID)
	if !ok || tracked.Response == nil || !tracked.DeleteWithInvocation || f.client == nil {
		return
	}
	_ = f.client.Rest.DeleteMessage(tracked.Response.ChannelID, tracked.Response.ID)
}