discord/bot.go

package discord

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"

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

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

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) Start(wg *sync.WaitGroup) error {
	defer wg.Done()

	err := bot.session.Open()
	if err != nil {
		return err
	}
	logrus.Info("Discord bot has connected")
	bot.running = make(chan os.Signal, 1)

	signal.Notify(bot.running, syscall.SIGINT, os.Kill)
	<-bot.running

	bot.session.Close()
	bot.online = false
	logrus.Info("Discord bot has been disconnected")

	return nil
}

func (bot *Bot) Stop() {
	if bot.online {
		logrus.Infof("Shutting down Bot")
		bot.running <- syscall.SIGINT
	}
}

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

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

	if err != nil {
		return bot, err
	}

	return bot, nil
}