arikawa/bot/ctx.go

package bot

import (
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"

	"github.com/diamondburned/arikawa/gateway"
	"github.com/diamondburned/arikawa/state"
	"github.com/pkg/errors"
)

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

// TODO: add variadic arguments

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

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

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

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

	s, err := state.New("Bot " + token)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create a dgo session")
	}

	c, err := New(s, cmd)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create rfrouter")
	}

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

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

	cancel := c.Start()

	if err := s.Open(); err != nil {
		return nil, errors.Wrap(err, "Failed to connect to Discord")
	}

	return func() error {
		Wait()
		// remove handler first
		cancel()
		// then finish closing session
		return s.Close()
	}, nil
}

// Wait blocks until SIGINT.
func Wait() {
	sigs := make(chan os.Signal)
	signal.Notify(sigs, os.Interrupt)
	<-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 := rfrouter.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,
		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
}

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

// FindCommand finds a command based on the struct and method name. The queried
// names will have their flags stripped.
//
// Example
//
//    // Find a command from the main context:
//    cmd := ctx.FindCommand("", "Method")
//    // Find a command from a subcommand:
//    cmd  = ctx.FindCommand("Starboard", "Reset")
//
func (ctx *Context) FindCommand(structname, methodname string) *CommandContext {
	if structname == "" {
		for _, c := range ctx.Commands {
			if c.MethodName == methodname {
				return c
			}
		}

		return nil
	}

	for _, sub := range ctx.subcommands {
		if sub.StructName != structname {
			continue
		}

		for _, c := range sub.Commands {
			if c.MethodName == methodname {
				return c
			}
		}
	}

	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.
func (ctx *Context) MustRegisterSubcommand(cmd interface{}) *Subcommand {
	s, err := ctx.RegisterSubcommand(cmd)
	if err != nil {
		panic(err)
	}

	return s
}

// RegisterSubcommand registers and adds cmd to the list of subcommands. It will
// also return the resulting Subcommand.
func (ctx *Context) RegisterSubcommand(cmd interface{}) (*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 err := s.InitCommands(ctx); err != nil {
		return nil, errors.Wrap(err, "Failed to initialize subcommand")
	}

	// Do a collision check
	for _, sub := range ctx.subcommands {
		if sub.Command == s.Command {
			return nil, errors.New(
				"New subcommand has duplicate name: " + s.Command)
		}
	}

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

// Start adds itself into the discordgo Session handlers. This needs to be run.
// The returned function is a delete function, which removes itself from the
// Session handlers.
func (ctx *Context) Start() func() {
	return ctx.State.AddHandler(func(v interface{}) {
		err := ctx.callCmd(v)
		if err == nil {
			return
		}

		str := ctx.FormatError(err)
		if str == "" {
			return
		}

		mc, isMessage := v.(*gateway.MessageCreateEvent)

		// Log the main error if reply is disabled or if the event isn't a
		// message.
		if !ctx.ReplyError || !isMessage {
			// Ignore trivial errors:
			switch err.(type) {
			case *ErrInvalidUsage, *ErrUnknownCommand:
				// Ignore
			default:
				ctx.ErrorLogger(errors.Wrap(err, "Command error"))
			}

			return
		}

		// Only reply if the event is not a message.
		if !isMessage {
			return
		}

		// Escape the error using the message sanitizer:
		str = ctx.SanitizeMessage(str)

		_, err = ctx.SendMessage(mc.ChannelID, str, nil)
		if err != nil {
			ctx.ErrorLogger(err)

			// TODO: there ought to be a better way lol
		}
	})
}

// 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 one. This function is used more for reference than an actual
// help message. As such, it only uses exported fields or methods.
func (ctx *Context) Help() string {
	return ctx.help(true)
}

func (ctx *Context) HelpAdmin() string {
	return ctx.help(false)
}

func (ctx *Context) help(hideAdmin bool) string {
	const indent = "      "

	var help strings.Builder

	// Generate the headers and descriptions
	help.WriteString("__Help__")

	if ctx.Name != "" {
		help.WriteString(": " + ctx.Name)
	}

	if ctx.Description != "" {
		help.WriteString("\n" + indent + ctx.Description)
	}

	if ctx.Flag.Is(AdminOnly) {
		// That's it.
		return help.String()
	}

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

	// Generate all commands
	help.WriteString("__Commands__")
	help.WriteString(ctx.Subcommand.Help(indent, hideAdmin))
	help.WriteByte('\n')

	var subHelp = strings.Builder{}
	var subcommands = ctx.Subcommands()

	for _, sub := range subcommands {
		if help := sub.Help(indent, hideAdmin); help != "" {
			for _, line := range strings.Split(help, "\n") {
				subHelp.WriteString(indent)
				subHelp.WriteString(line)
				subHelp.WriteByte('\n')
			}
		}
	}

	if subHelp.Len() > 0 {
		help.WriteString("---\n")
		help.WriteString("__Subcommands__\n")
		help.WriteString(subHelp.String())
	}

	return help.String()
}