discord-bot/bot.go

package discordbot

import (
	"fmt"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/sirupsen/logrus"
)

type Bot struct {
	token       string
	session     *discordgo.Session
	online      bool
	running     chan os.Signal
	slashRouter SlashRouter
}

var logHandler *logrus.Logger

func (bot *Bot) AddHandler(handler interface{}) error {
	// If the bot is online, we can not add a handler
	if bot.online {
		err := fmt.Errorf("unable to add handler to a bot that is already online")
		return err
	}

	bot.session.AddHandler(handler)
	return nil
}

func (bot *Bot) SetLegacyRouter(router *LegacyRouter) {
	bot.AddHandler(router.OnMessageCreate)
}

func (bot *Bot) SetSlashEventRouter(slashRouter *SlashRouter) {
	bot.slashRouter = *slashRouter
}

func (bot *Bot) Start(wg *sync.WaitGroup) error {
	err := bot.session.Open()
	if err != nil {
		return err
	}

	err = bot.slashRouter.GenerateCommands(bot.session)
	if err != nil {
		return err
	}

	logHandler.Info("The bot has started")
	bot.running = make(chan os.Signal, 1)
	bot.online = true

	go func() {
		for {
			select {
			case <-bot.running:
				logHandler.Info("Shutting down the bot")
				bot.slashRouter.CleanCommands(bot.session)
				bot.session.Close()
				logHandler.Info("The bot has gone offline")
				bot.online = false
				wg.Done()
			default:
				time.Sleep(time.Second * 1)
			}
		}
	}()

	return nil
}

func (bot *Bot) Stop() {
	if bot.online {
		bot.running <- syscall.SIGINT
	}
}

func NewBot(token string, log *logrus.Logger) (*Bot, error) {
	bot := &Bot{}
	bot.token = token

	logHandler = log

	var err error
	bot.session, err = discordgo.New("Bot " + bot.token)

	if err != nil {
		return bot, err
	}

	return bot, nil
}