arikawa/utils/bot/ctx.go

package bot

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/pkg/errors"

	"github.com/diamondburned/arikawa/v3/api"
	"github.com/diamondburned/arikawa/v3/utils/bot/extras/shellwords"
	"github.com/diamondburned/arikawa/v3/gateway"
	"github.com/diamondburned/arikawa/v3/gateway/shard"
	"github.com/diamondburned/arikawa/v3/session"
	"github.com/diamondburned/arikawa/v3/state"
	"github.com/diamondburned/arikawa/v3/state/store"
	"github.com/diamondburned/arikawa/v3/state/store/defaultstore"
)

// Prefixer checks a message if it starts with the desired prefix. By default,
// NewPrefix() is used.
type Prefixer func(*gateway.MessageCreateEvent) (prefix string, ok bool)

// NewPrefix creates a simple prefix checker using strings. As the default
// prefix is "!", the function is called as NewPrefix("!").
func NewPrefix(prefixes ...string) Prefixer {
	return func(msg *gateway.MessageCreateEvent) (string, bool) {
		for _, prefix := range prefixes {
			if strings.HasPrefix(msg.Content, prefix) {
				return prefix, true
			}
		}
		return "", false
	}
}

// ArgsParser is the function type for parsing message content into fields,
// usually delimited by spaces.
type ArgsParser func(content string) ([]string, error)

// DefaultArgsParser implements a parser similar to that of shell's,
// implementing quotes as well as escapes.
var DefaultArgsParser = shellwords.Parse

// NewShardFunc creates a shard constructor that shares the same internal store.
// If opts sets its own cabinet, then a new store isn't created.
func NewShardFunc(fn func(*state.State) (*Context, error)) shard.NewShardFunc {
	if fn == nil {
		panic("bot.NewShardFunc missing fn")
	}

	var once sync.Once
	var cab *store.Cabinet

	return func(m *shard.Manager, id *gateway.Identifier) (shard.Shard, error) {
		state := state.NewFromSession(session.NewCustomShard(m, id), nil)

		bot, err := fn(state)
		if err != nil {
			return nil, errors.Wrap(err, "failed to create bot instance")
		}

		if state.Cabinet == nil {
			once.Do(func() { cab = defaultstore.New() })
			state.Cabinet = cab
		}

		return bot, nil
	}
}

// Context is the bot state for commands and subcommands.
//
// Commands
//
// A command can be created by making it a method of Commands, or whatever
// struct was given to the constructor. This following example creates a command
// with a single integer argument (which can be ran with "~example 123"):
//
//    func (c *Commands) Example(
//        m *gateway.MessageCreateEvent, i int) (string, error) {
//
//        return fmt.Sprintf("You sent: %d", i)
//    }
//
// Commands' exported methods will all be used as commands. Messages are parsed
// with its first argument (the command) mapped accordingly to c.MapName, which
// capitalizes the first letter automatically to reflect the exported method
// name.
//
// A command can either return either an error, or data and error. The only data
// types allowed are string, *discord.Embed, and *api.SendMessageData. Any other
// return types will invalidate the method.
//
// Events
//
// An event can only have one argument, which is the pointer to the event
// struct. It can also only return error.
//
//    func (c *Commands) Example(o *gateway.TypingStartEvent) error {
//        log.Println("Someone's typing!")
//        return nil
//    }
type Context struct {
	*Subcommand
	*state.State

	// Descriptive (but optional) bot name
	Name string

	// Descriptive help body
	Description string

	// Called to parse message content, default to DefaultArgsParser().
	ParseArgs ArgsParser

	// Called to check a message's prefix. The default prefix is "!". Refer to
	// NewPrefix().
	HasPrefix Prefixer

	// AllowBot makes the router also process MessageCreate events from bots.
	// This is false by default and only applies to MessageCreate.
	AllowBot bool

	// QuietUnknownCommand, if true, will not make the bot reply with an unknown
	// command error into the chat. This will apply to all other subcommands.
	// SilentUnknown controls whether or not an UnknownCommandError should be
	// returned (instead of a silent error).
	SilentUnknown struct {
		// Command when true will silent only unknown commands. Known
		// subcommands with unknown commands will still error out.
		Command bool
		// Subcommand when true will suppress unknown subcommands.
		Subcommand bool
	}

	// FormatError formats any errors returned by anything, including the method
	// commands or the reflect functions. This also includes invalid usage
	// errors or unknown command errors. Returning an empty string means
	// ignoring the error.
	//
	// By default, this field replaces all @ with @\u200b, which prevents an
	// @everyone mention.
	FormatError func(error) string

	// ErrorLogger logs any error that anything makes and the library can't
	// reply to the client. This includes any event callback errors that aren't
	// Message Create.
	ErrorLogger func(error)

	// ReplyError when true replies to the user the error. This only applies to
	// MessageCreate events.
	ReplyError bool

	// ErrorReplier is an optional function that allows changing how the error
	// is replied. It overrides ReplyError and is only used for MessageCreate
	// events.
	//
	// Note that errors that are passed in here will bypas FormatError; in other
	// words, the implementation might only care about ErrorReplier and leave
	// FormatError as it is.
	ErrorReplier func(err error, src *gateway.MessageCreateEvent) api.SendMessageData

	// 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

	// Quick access map from event types to pointers. This map will never have
	// MessageCreateEvent's type.
	typeCache sync.Map // map[reflect.Type][]*CommandContext

	stopFunc func() // unbind function, see Start()
}

// Start quickly starts a bot with the given command. It will prepend "Bot"
// into the token automatically. Refer to example/ for usage.
func Start(
	token string, cmd interface{},
	opts func(*Context) error) (wait func() error, err error) {

	if token == "" {
		return nil, errors.New("token is not given")
	}

	if !strings.HasPrefix(token, "Bot ") {
		token = "Bot " + token
	}

	newShard := NewShardFunc(func(s *state.State) (*Context, error) {
		ctx, err := New(s, cmd)
		if err != nil {
			return nil, err
		}

		// fail api request if they (will) take up more than 5 minutes
		ctx.Client.Client.Timeout = 5 * time.Minute

		ctx.Gateway.ErrorLog = func(err error) {
			ctx.ErrorLogger(err)
		}

		if opts != nil {
			if err := opts(ctx); err != nil {
				return nil, err
			}
		}

		ctx.AddIntents(ctx.DeriveIntents())
		ctx.AddIntents(gateway.IntentGuilds) // for channel event caching

		return ctx, nil
	})

	m, err := shard.NewManager(token, newShard)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create shard manager")
	}

	if err := m.Open(context.Background()); err != nil {
		return nil, errors.Wrap(err, "failed to open")
	}

	return func() error {
		WaitForInterrupt()

		// Close the shards first.
		closeErr := m.Close()

		// Remove all handlers to clean up.
		m.ForEach(func(s shard.Shard) {
			ctx := s.(*Context)

			stop := ctx.Start()
			stop()
		})

		return closeErr
	}, nil
}

// Run starts the bot, prints a message into the console, and blocks until
// SIGINT. "Bot" is prepended into the token automatically, similar to Start.
// The function will call os.Exit(1) on an initialization or cleanup error.
func Run(token string, cmd interface{}, opts func(*Context) error) {
	wait, err := Start(token, cmd, opts)
	if err != nil {
		log.Fatalln("failed to start:", err)
	}

	log.Println("Bot is running.")

	if err := wait(); err != nil {
		log.Fatalln("cleanup error:", err)
	}
}

// WaitForInterrupt blocks until SIGINT.
func WaitForInterrupt() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, syscall.SIGTERM)
	<-sigs
}

// New makes a new context with a "~" as the prefix. cmds must be a pointer to a
// struct with a *Context field. Example:
//
//    type Commands struct {
//        Ctx *Context
//    }
//
//    cmds := &Commands{}
//    c, err := bot.New(session, cmds)
//
// The default prefix is "~", which means commands must start with "~" followed
// by the command name in the first argument, else it will be ignored.
//
// c.Start() should be called afterwards to actually handle incoming events.
func New(s *state.State, cmd interface{}) (*Context, error) {
	c, err := NewSubcommand(cmd)
	if err != nil {
		return nil, err
	}

	ctx := &Context{
		Subcommand: c,
		State:      s,
		ParseArgs:  DefaultArgsParser,
		HasPrefix:  NewPrefix("~"),
		FormatError: func(err error) string {
			// Escape all pings, including @everyone.
			return strings.Replace(err.Error(), "@", "@\u200b", -1)
		},
		ErrorLogger: func(err error) {
			log.Println("Bot error:", err)
		},
		ReplyError: true,
	}

	if err := ctx.InitCommands(ctx); err != nil {
		return nil, errors.Wrap(err, "failed to initialize with given cmds")
	}

	return ctx, nil
}

// AddIntents adds the given Gateway Intent into the Gateway. This is a
// convenient function that calls Gateway's AddIntent.
func (ctx *Context) AddIntents(i gateway.Intents) {
	ctx.Gateway.AddIntents(i)
}

// Subcommands returns the slice of subcommands. To add subcommands, use
// RegisterSubcommand().
func (ctx *Context) Subcommands() []*Subcommand {
	// Getter is not useless, refer to the struct doc for reason.
	return ctx.subcommands
}

// FindMethod finds a method based on the struct and method name. The queried
// names will have their flags stripped.
//
//    // Find a command from the main context:
//    cmd := ctx.FindMethod("", "Method")
//    // Find a command from a subcommand:
//    cmd  = ctx.FindMethod("Starboard", "Reset")
//
func (ctx *Context) FindCommand(structName, methodName string) *MethodContext {
	if structName == "" {
		return ctx.Subcommand.FindCommand(methodName)
	}
	for _, sub := range ctx.subcommands {
		if sub.StructName == structName {
			return sub.FindCommand(methodName)
		}
	}
	return nil
}

// MustRegisterSubcommand tries to register a subcommand, and will panic if it
// fails. This is recommended, as subcommands won't change after initializing
// once in runtime, thus fairly harmless after development.
//
// If no names are given or if the first name is empty, then the subcommand name
// will be derived from the struct name. If one name is given, then that name
// will override the struct name. Any other name values will be aliases.
//
// It is recommended to use this method to add subcommand aliases over manually
// altering the Aliases slice of each Subcommand, as it does collision checks
// against other subcommands as well.
func (ctx *Context) MustRegisterSubcommand(cmd interface{}, names ...string) *Subcommand {
	s, err := ctx.RegisterSubcommand(cmd, names...)
	if err != nil {
		panic(err)
	}
	return s
}

// RegisterSubcommand registers and adds cmd to the list of subcommands. It will
// also return the resulting Subcommand. Refer to MustRegisterSubcommand for the
// names argument.
func (ctx *Context) RegisterSubcommand(cmd interface{}, names ...string) (*Subcommand, error) {
	s, err := NewSubcommand(cmd)
	if err != nil {
		return nil, errors.Wrap(err, "failed to add subcommand")
	}

	// Register the subcommand's name.
	s.NeedsName()

	if len(names) > 0 && names[0] != "" {
		s.Command = names[0]
	}

	if len(names) > 1 {
		// Copy the slice for expected behaviors.
		s.Aliases = append([]string(nil), names[1:]...)
	}

	if err := s.InitCommands(ctx); err != nil {
		return nil, errors.Wrap(err, "failed to initialize subcommand")
	}

	// Check if the existing command name already exists. This could really be
	// optimized, but since it's in a cold path, who cares.
	var subcommandNames = append([]string{s.Command}, s.Aliases...)

	for _, name := range subcommandNames {
		for _, sub := range ctx.subcommands {
			// Check each alias against the subcommand name.
			if sub.Command == name {
				return nil, fmt.Errorf("new subcommand has duplicate name: %q", name)
			}

			// Also check each alias against other subcommands' aliases.
			for _, subalias := range sub.Aliases {
				if subalias == name {
					return nil, fmt.Errorf("new subcommand has duplicate alias: %q", name)
				}
			}
		}
	}

	ctx.subcommands = append(ctx.subcommands, s)
	return s, nil
}

// emptyMentionTypes is used by Start() to not parse any mentions.
var emptyMentionTypes = []api.AllowedMentionType{}

// Start adds itself into the session handlers. If Start is called more than
// once, then it does nothing. The caller doesn't have to call Start if they
// call Open.
//
// The returned function is a delete function, which removes itself from the
// Session handlers. The delete function is not safe to use concurrently.
func (ctx *Context) Start() func() {
	if ctx.stopFunc == nil {
		cancel := ctx.State.AddHandler(func(v interface{}) {
			if err := ctx.callCmd(v); err != nil {
				ctx.ErrorLogger(errors.Wrap(err, "command error"))
			}
		})

		ctx.stopFunc = func() {
			cancel()
			ctx.stopFunc = nil
		}
	}

	return ctx.stopFunc
}

// Open starts the bot context and the gateway connection. It automatically
// binds the needed handlers.
func (ctx *Context) Open(cancelCtx context.Context) error {
	ctx.Start()
	return ctx.State.Open(cancelCtx)
}

// Call should only be used if you know what you're doing.
func (ctx *Context) Call(event interface{}) error {
	return ctx.callCmd(event)
}

// Help generates a full Help message. It serves mainly as a reference for
// people to reimplement and change. It doesn't show hidden commands.
func (ctx *Context) Help() string {
	return ctx.HelpGenerate(false)
}

// HelpGenerate generates a full Help message. It serves mainly as a reference
// for people to reimplement and change. If showHidden is true, then hidden
// subcommands and commands will be shown.
func (ctx *Context) HelpGenerate(showHidden bool) string {
	// Generate the header.
	buf := strings.Builder{}
	buf.WriteString("__Help__")

	// Name an
	if ctx.Name != "" {
		buf.WriteString(": " + ctx.Name)
	}
	if ctx.Description != "" {
		buf.WriteString("\n" + IndentLines(ctx.Description))
	}

	// Separators
	buf.WriteString("\n---\n")

	// Generate all commands
	if help := ctx.Subcommand.Help(); help != "" {
		buf.WriteString("__Commands__\n")
		buf.WriteString(IndentLines(help))
		buf.WriteByte('\n')
	}

	var subcommands = ctx.Subcommands()
	var subhelps = make([]string, 0, len(subcommands))

	for _, sub := range subcommands {
		if sub.Hidden && !showHidden {
			continue
		}

		help := sub.HelpShowHidden(showHidden)
		if help == "" {
			continue
		}

		help = IndentLines(help)

		builder := strings.Builder{}
		builder.WriteString("**")
		builder.WriteString(sub.Command)
		builder.WriteString("**")

		for _, alias := range sub.Aliases {
			builder.WriteString("|")
			builder.WriteString("**")
			builder.WriteString(alias)
			builder.WriteString("**")
		}

		if sub.Description != "" {
			builder.WriteString(": ")
			builder.WriteString(sub.Description)
		}

		builder.WriteByte('\n')
		builder.WriteString(help)

		subhelps = append(subhelps, builder.String())
	}

	if len(subhelps) > 0 {
		buf.WriteString("---\n")
		buf.WriteString("__Subcommands__\n")
		buf.WriteString(IndentLines(strings.Join(subhelps, "\n")))
	}

	return buf.String()
}

// IndentLine prefixes every line from input with a single-level indentation.
func IndentLines(input string) string {
	const indent = "      "
	var lines = strings.Split(input, "\n")
	for i := range lines {
		lines[i] = indent + lines[i]
	}
	return strings.Join(lines, "\n")
}

// DeriveIntents derives all possible gateway intents from this context and all
// its subcommands' method handlers and middlewares.
func (ctx *Context) DeriveIntents() gateway.Intents {
	var intents = ctx.Subcommand.DeriveIntents()
	for _, subcmd := range ctx.subcommands {
		intents |= subcmd.DeriveIntents()
	}
	return intents
}