unbot/features/iconchanger.go

package features

import (
	"io"
	"log/slog"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/disgoorg/disgo/bot"
	"github.com/disgoorg/disgo/discord"
	"github.com/disgoorg/disgo/events"
	"github.com/disgoorg/omit"
	"github.com/disgoorg/snowflake/v2"
)

type IconChanger struct {
	client    *bot.Client
	guildID   snowflake.ID
	channelID snowflake.ID
	interval  time.Duration
	enabled   bool
	mu        sync.RWMutex
	imageURLs []string
	stop      chan struct{}
}

func NewIconChanger(client *bot.Client, guildID, channelID uint64, interval time.Duration, enabled bool) *IconChanger {
	return &IconChanger{
		client:    client,
		guildID:   snowflake.ID(guildID),
		channelID: snowflake.ID(channelID),
		interval:  interval,
		enabled:   enabled,
		stop:      make(chan struct{}),
	}
}

func (ic *IconChanger) HandleMessage(e *events.MessageCreate) {
	if !ic.enabled || e.ChannelID != ic.channelID || e.Message.Author.Bot {
		return
	}

	for _, a := range e.Message.Attachments {
		if a.ContentType != nil && strings.HasPrefix(*a.ContentType, "image/") {
			ic.mu.Lock()
			ic.imageURLs = append(ic.imageURLs, a.URL)
			ic.mu.Unlock()
		}
	}

	for _, embed := range e.Message.Embeds {
		if embed.Image != nil {
			ic.mu.Lock()
			ic.imageURLs = append(ic.imageURLs, embed.Image.URL)
			ic.mu.Unlock()
		}
	}
}

func (ic *IconChanger) Start() {
	if !ic.enabled {
		return
	}

	interval := ic.interval
	if interval <= 0 {
		interval = 1 * time.Hour
	}

	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				ic.changeIcon()
			case <-ic.stop:
				return
			}
		}
	}()
}

func (ic *IconChanger) Stop() {
	close(ic.stop)
}

func (ic *IconChanger) BackfillHistory() {
	if !ic.enabled {
		return
	}

	slog.Info("starting icon history backfill", slog.String("channel_id", ic.channelID.String()))

	var lastID snowflake.ID
	seen := make(map[string]struct{})

	for {
		messages, err := ic.client.Rest.GetMessages(ic.channelID, 0, lastID, 0, 100)
		if err != nil {
			slog.Error("failed to get messages for icon backfill", slog.Any("err", err))
			return
		}
		if len(messages) == 0 {
			break
		}

		for _, msg := range messages {
			if msg.Author.Bot {
				continue
			}

			for _, a := range msg.Attachments {
				if a.ContentType != nil && strings.HasPrefix(*a.ContentType, "image/") {
					if _, ok := seen[a.URL]; !ok {
						seen[a.URL] = struct{}{}
						ic.mu.Lock()
						ic.imageURLs = append(ic.imageURLs, a.URL)
						ic.mu.Unlock()
					}
				}
			}

			for _, embed := range msg.Embeds {
				if embed.Image != nil {
					if _, ok := seen[embed.Image.URL]; !ok {
						seen[embed.Image.URL] = struct{}{}
						ic.mu.Lock()
						ic.imageURLs = append(ic.imageURLs, embed.Image.URL)
						ic.mu.Unlock()
					}
				}
			}
		}

		lastID = messages[len(messages)-1].ID

		time.Sleep(200 * time.Millisecond)
	}

	slog.Info("icon history backfill complete", slog.Int("images_found", len(ic.imageURLs)))
}

func (ic *IconChanger) changeIcon() {
	ic.mu.RLock()
	urls := make([]string, len(ic.imageURLs))
	copy(urls, ic.imageURLs)
	ic.mu.RUnlock()

	if len(urls) == 0 {
		return
	}

	url := urls[rand.Intn(len(urls))]
	icon, err := downloadIcon(url)
	if err != nil {
		slog.Error("failed to download icon", slog.Any("err", err))
		return
	}

	_, err = ic.client.Rest.UpdateGuild(ic.guildID, discord.GuildUpdate{
		Icon: omit.New(icon),
	})
	if err != nil {
		slog.Error("failed to update guild icon", slog.Any("err", err))
		return
	}

	slog.Info("guild icon changed", slog.String("url", url))
}

func downloadIcon(url string) (*discord.Icon, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	mediaType := resp.Header.Get("Content-Type")
	if mediaType == "" {
		mediaType = "image/png"
	}

	var iconType discord.IconType
	switch {
	case mediaType == "image/png":
		iconType = discord.IconTypePNG
	case mediaType == "image/jpeg":
		iconType = discord.IconTypeJPEG
	case mediaType == "image/webp":
		iconType = discord.IconTypeWEBP
	case mediaType == "image/gif":
		iconType = discord.IconTypeGIF
	default:
		iconType = discord.IconTypePNG
	}

	return discord.NewIconRaw(iconType, data), nil
}