arikawa/state/store_default.go

package state

import (
	"sort"
	"sync"

	"github.com/diamondburned/arikawa/discord"
)

// TODO: make an ExpiryStore

type DefaultStore struct {
	*DefaultStoreOptions

	self discord.User

	// includes normal and private
	privates map[discord.Snowflake]*discord.Channel  // channelID:channel
	guilds   map[discord.Snowflake]*discord.Guild    // guildID:guild
	messages map[discord.Snowflake][]discord.Message // channelID:messages

	mut sync.Mutex
}

type DefaultStoreOptions struct {
	MaxMessages uint // default 50
}

var _ Store = (*DefaultStore)(nil)

func NewDefaultStore(opts *DefaultStoreOptions) *DefaultStore {
	if opts == nil {
		opts = &DefaultStoreOptions{
			MaxMessages: 50,
		}
	}

	return &DefaultStore{
		DefaultStoreOptions: opts,

		privates: map[discord.Snowflake]*discord.Channel{},
		guilds:   map[discord.Snowflake]*discord.Guild{},
		messages: map[discord.Snowflake][]discord.Message{},
	}
}

func (s *DefaultStore) Reset() error {
	s.mut.Lock()
	defer s.mut.Unlock()

	s.self = discord.User{}

	s.privates = map[discord.Snowflake]*discord.Channel{}
	s.guilds = map[discord.Snowflake]*discord.Guild{}
	s.messages = map[discord.Snowflake][]discord.Message{}

	return nil
}

////

func (s *DefaultStore) Self() (*discord.User, error) {
	s.mut.Lock()
	defer s.mut.Unlock()

	if !s.self.ID.Valid() {
		return nil, ErrStoreNotFound
	}

	return &s.self, nil
}

func (s *DefaultStore) SetSelf(me *discord.User) error {
	s.mut.Lock()
	s.self = *me
	s.mut.Unlock()

	return nil
}

////

func (s *DefaultStore) Channel(id discord.Snowflake) (*discord.Channel, error) {
	s.mut.Lock()
	defer s.mut.Unlock()

	for _, g := range s.guilds {
		for _, ch := range g.Channels {
			if ch.ID == id {
				return &ch, nil
			}
		}
	}

	return nil, ErrStoreNotFound
}

func (s *DefaultStore) Channels(
	guildID discord.Snowflake) ([]discord.Channel, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return gd.Channels, nil
}

func (s *DefaultStore) PrivateChannels() ([]discord.Channel, error) {
	s.mut.Lock()

	var chs = make([]discord.Channel, 0, len(s.privates))
	for _, ch := range s.privates {
		chs = append(chs, *ch)
	}

	s.mut.Unlock()

	sort.Slice(chs, func(i, j int) bool {
		// Latest first
		return chs[i].LastMessageID > chs[j].LastMessageID
	})

	return chs, nil
}

func (s *DefaultStore) ChannelSet(channel *discord.Channel) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	switch channel.Type {
	case discord.DirectMessage, discord.GroupDM:
		s.privates[channel.ID] = channel

	default:
		gd, ok := s.guilds[channel.GuildID]
		if !ok {
			return ErrStoreNotFound
		}

		for i, ch := range gd.Channels {
			if ch.ID == channel.ID {
				// Found, just edit
				gd.Channels[i] = *channel
				return nil
			}
		}

		gd.Channels = append(gd.Channels, *channel)
	}

	return nil
}

func (s *DefaultStore) ChannelRemove(channel *discord.Channel) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[channel.GuildID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, ch := range gd.Channels {
		if ch.ID == channel.ID {
			gd.Channels = append(gd.Channels[:i], gd.Channels[i+1:]...)
			return nil
		}
	}

	return ErrStoreNotFound
}

////

func (s *DefaultStore) Emoji(
	guildID, emojiID discord.Snowflake) (*discord.Emoji, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	for _, emoji := range gd.Emojis {
		if emoji.ID == emojiID {
			return &emoji, nil
		}
	}

	return nil, ErrStoreNotFound
}

func (s *DefaultStore) Emojis(
	guildID discord.Snowflake) ([]discord.Emoji, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return gd.Emojis, nil
}

func (s *DefaultStore) EmojiSet(
	guildID discord.Snowflake, emojis []discord.Emoji) error {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	filtered := emojis[:0]

Main:
	for _, enew := range emojis {
		// Try and see if this emoji is already in the slice
		for i, emoji := range gd.Emojis {
			if emoji.ID == enew.ID {
				// If it is, we simply replace it
				gd.Emojis[i] = enew

				continue Main
			}
		}

		// If not, we add it to the slice that's to be appended.
		filtered = append(filtered, enew)
	}

	// Append the new emojis
	gd.Emojis = append(gd.Emojis, filtered...)
	return nil
}

////

func (s *DefaultStore) Guild(id discord.Snowflake) (*discord.Guild, error) {
	s.mut.Lock()
	defer s.mut.Unlock()

	ch, ok := s.guilds[id]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return ch, nil
}

func (s *DefaultStore) Guilds() ([]discord.Guild, error) {
	s.mut.Lock()

	var gs = make([]discord.Guild, 0, len(s.guilds))
	for _, g := range s.guilds {
		gs = append(gs, *g)
	}

	s.mut.Unlock()

	sort.Slice(gs, func(i, j int) bool {
		return gs[i].ID > gs[j].ID
	})

	return gs, nil
}

func (s *DefaultStore) GuildSet(g *discord.Guild) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	old := s.guilds[g.ID]

	// Check the channels too
	if len(old.Channels) == 0 {
		return nil
	}

	if len(g.Channels) == 0 {
		g.Channels = old.Channels
	}

	s.guilds[g.ID] = g
	return nil
}

func (s *DefaultStore) GuildRemove(g *discord.Guild) error {
	s.mut.Lock()
	delete(s.guilds, g.ID)
	s.mut.Unlock()

	return nil
}

////

func (s *DefaultStore) Member(
	guildID, userID discord.Snowflake) (*discord.Member, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	for _, member := range gd.Members {
		if member.User.ID == userID {
			return &member, nil
		}
	}

	return nil, ErrStoreNotFound
}

func (s *DefaultStore) Members(
	guildID discord.Snowflake) ([]discord.Member, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return gd.Members, nil
}

func (s *DefaultStore) MemberSet(
	guildID discord.Snowflake, member *discord.Member) error {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	// Try and see if this member is already in the slice
	for i, m := range gd.Members {
		if m.User.ID == member.User.ID {
			// If it is, we simply replace it
			gd.Members[i] = *member
			return nil
		}
	}

	// Append the new member
	gd.Members = append(gd.Members, *member)
	return nil
}

func (s *DefaultStore) MemberRemove(guildID, userID discord.Snowflake) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	// Try and see if this member is already in the slice
	for i, m := range gd.Members {
		if m.User.ID == userID {
			gd.Members = append(gd.Members[:i], gd.Members[i+1:]...)
			return nil
		}
	}

	return ErrStoreNotFound
}

////

func (s *DefaultStore) Message(
	channelID, messageID discord.Snowflake) (*discord.Message, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	ms, ok := s.messages[channelID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	for _, m := range ms {
		if m.ID == messageID {
			return &m, nil
		}
	}

	return nil, ErrStoreNotFound
}

func (s *DefaultStore) Messages(
	channelID discord.Snowflake) ([]discord.Message, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	ms, ok := s.messages[channelID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return ms, nil
}

func (s *DefaultStore) MessageSet(message *discord.Message) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	ms, ok := s.messages[message.ChannelID]
	if !ok {
		ms = make([]discord.Message, 0, int(s.MaxMessages)+1)
	}

	// Append
	ms = append(ms, *message)

	// Sort (should be fast since it's presorted)
	sort.Slice(ms, func(i, j int) bool {
		return ms[i].ID > ms[j].ID
	})

	if len(ms) > int(s.MaxMessages) {
		ms = ms[len(ms)-int(s.MaxMessages):]
	}

	return nil
}

func (s *DefaultStore) MessageRemove(message *discord.Message) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	ms, ok := s.messages[message.ChannelID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, m := range ms {
		if m.ID == message.ID {
			ms = append(ms[:i], ms[i+1:]...)
			s.messages[message.ChannelID] = ms
			return nil
		}
	}

	return ErrStoreNotFound
}

////

func (s *DefaultStore) Presence(
	guildID, userID discord.Snowflake) (*discord.Presence, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	for _, p := range gd.Presences {
		if p.User.ID == userID {
			return &p, nil
		}
	}

	return nil, ErrStoreNotFound
}

func (s *DefaultStore) Presences(
	guildID discord.Snowflake) ([]discord.Presence, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return gd.Presences, nil
}

func (s *DefaultStore) PresenceSet(
	guildID discord.Snowflake, presence *discord.Presence) error {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, p := range gd.Presences {
		if p.User.ID == presence.User.ID {
			gd.Presences[i] = *presence
			return nil
		}
	}

	gd.Presences = append(gd.Presences, *presence)
	return nil
}

func (s *DefaultStore) PresenceRemove(guildID, userID discord.Snowflake) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, p := range gd.Presences {
		if p.User.ID == userID {
			gd.Presences = append(gd.Presences[:i], gd.Presences[i+1:]...)
			return nil
		}
	}

	return ErrStoreNotFound
}

////

func (s *DefaultStore) Role(
	guildID, roleID discord.Snowflake) (*discord.Role, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	for _, r := range gd.Roles {
		if r.ID == roleID {
			return &r, nil
		}
	}

	return nil, ErrStoreNotFound
}

func (s *DefaultStore) Roles(
	guildID discord.Snowflake) ([]discord.Role, error) {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return nil, ErrStoreNotFound
	}

	return gd.Roles, nil
}

func (s *DefaultStore) RoleSet(
	guildID discord.Snowflake, role *discord.Role) error {

	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, r := range gd.Roles {
		if r.ID == role.ID {
			gd.Roles[i] = *role
			return nil
		}
	}

	gd.Roles = append(gd.Roles, *role)
	return nil
}

func (s *DefaultStore) RoleRemove(guildID, roleID discord.Snowflake) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	gd, ok := s.guilds[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, r := range gd.Roles {
		if r.ID == roleID {
			gd.Roles = append(gd.Roles[:i], gd.Roles[i+1:]...)
			return nil
		}
	}

	return ErrStoreNotFound
}