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

	channels  map[discord.Snowflake][]discord.Channel  // guildID:channels
	members   map[discord.Snowflake][]discord.Member   // guildID:members
	presences map[discord.Snowflake][]discord.Presence // guildID:presences
	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,
		}
	}

	ds := &DefaultStore{
		DefaultStoreOptions: opts,
	}
	ds.Reset()

	return ds
}

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.channels = map[discord.Snowflake][]discord.Channel{}
	s.members = map[discord.Snowflake][]discord.Member{}
	s.presences = map[discord.Snowflake][]discord.Presence{}
	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) SelfSet(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 _, chs := range s.channels {
		for _, ch := range chs {
			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()

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

	return append([]discord.Channel{}, chs...), 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:
		chs := s.channels[channel.GuildID]

		for i, ch := range chs {
			if ch.ID == channel.ID {
				// Also from discordgo.
				if channel.Permissions == nil {
					channel.Permissions = ch.Permissions
				}

				// Found, just edit
				chs[i] = *channel

				return nil
			}
		}

		chs = append(chs, *channel)
		s.channels[channel.GuildID] = chs
	}

	return nil
}

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

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

	for i, ch := range chs {
		if ch.ID == channel.ID {
			chs = append(chs[:i], chs[i+1:]...)
			s.channels[channel.GuildID] = chs

			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 append([]discord.Emoji{}, 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()

	if len(s.guilds) == 0 {
		s.mut.Unlock()
		return nil, ErrStoreNotFound
	}

	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(guild *discord.Guild) error {
	s.mut.Lock()
	defer s.mut.Unlock()

	if g, ok := s.guilds[guild.ID]; ok {
		// preserve state stuff
		if guild.Roles == nil {
			guild.Roles = g.Roles
		}
		if guild.Emojis == nil {
			guild.Emojis = g.Emojis
		}
	}

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

func (s *DefaultStore) GuildRemove(id discord.Snowflake) error {
	s.mut.Lock()
	delete(s.guilds, id)
	s.mut.Unlock()

	return nil
}

////

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

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

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

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

	return nil, ErrStoreNotFound
}

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

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

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

	return append([]discord.Member{}, ms...), nil
}

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

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

	ms := s.members[guildID]

	// Try and see if this member is already in the slice
	for i, m := range ms {
		if m.User.ID == member.User.ID {
			// If it is, we simply replace it
			ms[i] = *member
			s.members[guildID] = ms

			return nil
		}
	}

	// Append the new member
	ms = append(ms, *member)
	s.members[guildID] = ms

	return nil
}

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

	ms, ok := s.members[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	// Try and see if this member is already in the slice
	for i, m := range ms {
		if m.User.ID == userID {
			ms = append(ms, ms[i+1:]...)
			s.members[guildID] = ms

			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 append([]discord.Message{}, ms...), nil
}

func (s *DefaultStore) MaxMessages() int {
	return int(s.DefaultStoreOptions.MaxMessages)
}

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, s.MaxMessages()+1)
	}

	// Check if we already have the message.
	for i, m := range ms {
		if m.ID == message.ID {
			// Thanks, Discord.
			if message.Content != "" {
				m.Content = message.Content
			}
			if message.EditedTimestamp != nil {
				m.EditedTimestamp = message.EditedTimestamp
			}
			if message.Mentions != nil {
				m.Mentions = message.Mentions
			}
			if message.Embeds != nil {
				m.Embeds = message.Embeds
			}
			if message.Attachments != nil {
				m.Attachments = message.Attachments
			}
			if message.Timestamp.Valid() {
				m.Timestamp = message.Timestamp
			}
			if message.Author.ID.Valid() {
				m.Author = message.Author
			}

			ms[i] = m
			return nil
		}
	}

	// Prepend the latest message at the end

	if len(ms) > 0 {
		var end = s.MaxMessages()
		if len(ms) < end {
			end = len(ms)
		}

		// Copy hack to prepend. This copies the 0th-(end-1)th entries to
		// 1st-endth.
		copy(ms[1:end], ms[0:end-1])
		// Then, set the 0th entry.
		ms[0] = *message

	} else {
		ms = append(ms, *message)
	}

	s.messages[message.ChannelID] = ms
	return nil
}

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

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

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

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

	return ErrStoreNotFound
}

////

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

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

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

	for _, p := range ps {
		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()

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

	return append([]discord.Presence{}, ps...), nil
}

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

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

	ps := s.presences[guildID]

	for i, p := range ps {
		if p.User.ID == presence.User.ID {
			ps[i] = *presence
			s.presences[guildID] = ps

			return nil
		}
	}

	ps = append(ps, *presence)
	s.presences[guildID] = ps
	return nil
}

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

	ps, ok := s.presences[guildID]
	if !ok {
		return ErrStoreNotFound
	}

	for i, p := range ps {
		if p.User.ID == userID {
			ps = append(ps[:i], ps[i+1:]...)
			s.presences[guildID] = ps

			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 append([]discord.Role{}, 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
}