prun-trade-bot-go/lib/bot/bot.go

package bot

import (
	"code.lowsec.club/okawari/go-discord-bot/lib/commands"
	"context"
	"fmt"
	"github.com/bwmarrin/discordgo"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
)

type Bot struct {
	discord  *discordgo.Session
	logger   zerolog.Logger
	guildId  string
	commands []struct {
		Command       *commands.SlashCommand
		removeOnClose bool
	}
	registeredCmds []*discordgo.ApplicationCommand
	ctx            context.Context
}

func (b *Bot) AddCommand(command commands.SlashCommand, removeOnClose bool) *Bot {
	b.commands = append(b.commands, struct {
		Command       *commands.SlashCommand
		removeOnClose bool
	}{Command: &command, removeOnClose: removeOnClose})

	return b
}

func (b *Bot) AddCommands(commands []commands.SlashCommand, removeOnClose bool) *Bot {
	for _, cmd := range commands {
		b.AddCommand(cmd, removeOnClose)
	}

	return b
}

func (b *Bot) registerCommand(command commands.SlashCommand, removeOnClose bool) error {
	acc := command.CreateCommand()
	logger := b.logger.With().Str("cmd", acc.Name).Logger()

	registeredCmd, err := b.discord.ApplicationCommandCreate(b.discord.State.User.ID, b.guildId, acc)
	if err != nil {
		return err
	}
	logger.Debug().Msg("Command registered")

	if removeOnClose == true {
		b.registeredCmds = append(b.registeredCmds, registeredCmd)
	}

	b.discord.AddHandler(func(s *discordgo.Session, r *discordgo.InteractionCreate) {
		logger = logger.With().
			Str("user", r.Member.Nick).
			Str("component", "Handler").
			Logger()
		if r.ApplicationCommandData().ID == registeredCmd.ID {
			ctx := logger.WithContext(b.ctx)
			command.Handle(ctx, s, r)
		}
	})

	return nil
}

func (b *Bot) Open() error {
	b.discord.AddHandler(b.onReady())
	return b.discord.Open()
}

func (b *Bot) onReady() func(s *discordgo.Session, r *discordgo.Ready) {
	return func(s *discordgo.Session, r *discordgo.Ready) {
		log.Info().
			Str("user", fmt.Sprintf("%v#%v", s.State.User.Username, s.State.User.Discriminator)).
			Msg("Logged inn")

		log.Info().Int("commands", len(b.commands)).Msg("Registering commands")
		for _, cmd := range b.commands {
			err := b.registerCommand(*cmd.Command, cmd.removeOnClose)
			if err != nil {
				b.logger.Fatal().Err(err).Msg("Unable to register command")
			}
		}
	}
}

func (b *Bot) Close() {
	b.logger.Debug().Msg("Shutting down")

	b.logger.Debug().Int("commands", len(b.registeredCmds)).Msg("Registered commands")

	for _, cmd := range b.registeredCmds {
		log.Info().
			Str("name", cmd.Name).
			Msg("Unregistering command")

		err := b.discord.ApplicationCommandDelete(b.discord.State.User.ID, b.guildId, cmd.ID)
		if err != nil {
			log.Error().
				Err(err).
				Str("name", cmd.Name).Msg("Failed to unregister command")
		}
	}

	err := b.discord.Close()
	if err != nil {
		b.logger.Error().Err(err).Msg("Unable to close discord")
	}
}

func New(ctx context.Context, name string, token string) (*Bot, error) {

	discord, err := discordgo.New("Bot " + token)
	if err != nil {
		return nil, fmt.Errorf("unable to create discord bot: %s", err)
	}

	return (&Bot{
		discord: discord,
		logger:  log.With().Str("bot", name).Logger(),
		ctx:     ctx,
		guildId: "758249214160797706",
		commands: make([]struct {
			Command       *commands.SlashCommand
			removeOnClose bool
		}, 0),
		registeredCmds: make([]*discordgo.ApplicationCommand, 0),
	}), nil
}

func (b *Bot) init() *Bot {

	return b
}