velox-bot/internal/music/manager.go

package music

import (
	"context"
	"fmt"
	"net/url"
	"sync"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/disgoorg/disgolink/v3/disgolink"
	"github.com/disgoorg/disgolink/v3/lavalink"
	"github.com/disgoorg/snowflake/v2"
)

type TrackEntry struct {
	Track       lavalink.Track
	RequestedBy string
}

type guildPlayer struct {
	Player           disgolink.Player
	Queue            []TrackEntry
	RepeatSong       bool
	RepeatQueue      bool
	Paused           bool
	Volume           int
	IdleSince        time.Time
	TextChannelID    string
	NowPlayingMsgID  string
}

type Manager struct {
	client  disgolink.Client
	session *discordgo.Session

	mu      sync.Mutex
	players map[string]*guildPlayer
}

var manager *Manager

func Init(session *discordgo.Session, appID, lavalinkHost, lavalinkPass string) error {
	if lavalinkHost == "" {
		return nil
	}

	userID, err := snowflake.Parse(appID)
	if err != nil {
		return fmt.Errorf("parse app id for lavalink: %w", err)
	}

	m := &Manager{
		client:  disgolink.New(userID, disgolink.WithListenerFunc(onTrackEnd)),
		session: session,
		players: make(map[string]*guildPlayer),
	}

	u, err := url.Parse(lavalinkHost)
	if err != nil {
		return fmt.Errorf("parse lavalink host: %w", err)
	}

	secure := u.Scheme == "https" || u.Scheme == "wss"
	password := lavalinkPass
	if password == "" {
		password = "youshallnotpass"
	}

	_, err = m.client.AddNode(context.Background(), disgolink.NodeConfig{
		Name:     "main",
		Address:  u.Host,
		Password: password,
		Secure:   secure,
	})
	if err != nil {
		return fmt.Errorf("add lavalink node: %w", err)
	}

	manager = m
	go idleDisconnectLoop()
	return nil
}

func idleDisconnectLoop() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		if manager == nil || manager.session == nil {
			continue
		}

		now := time.Now()
		var toDisconnect []string

		manager.mu.Lock()
		for guildID, gp := range manager.players {
			if gp == nil {
				continue
			}

			// If we're paused or have anything queued/playing, clear idle timer.
			if gp.Paused || len(gp.Queue) > 0 {
				gp.IdleSince = time.Time{}
				continue
			}

			if gp.IdleSince.IsZero() {
				gp.IdleSince = now
				continue
			}

			if now.Sub(gp.IdleSince) >= time.Minute {
				toDisconnect = append(toDisconnect, guildID)
				gp.IdleSince = time.Time{}
			}
		}
		manager.mu.Unlock()

		for _, guildID := range toDisconnect {
			// Best-effort: disconnect. Safe even if we're already disconnected.
			_ = manager.session.ChannelVoiceJoinManual(guildID, "", false, false)
		}
	}
}

func OnVoiceStateUpdate(vs *discordgo.VoiceStateUpdate) {
	if manager == nil {
		return
	}
	if vs.UserID != manager.session.State.User.ID {
		return
	}

	guildID := snowflake.MustParse(vs.GuildID)
	var channelID *snowflake.ID
	if vs.ChannelID != "" {
		id := snowflake.MustParse(vs.ChannelID)
		channelID = &id
	}

	manager.client.OnVoiceStateUpdate(
		context.Background(),
		guildID,
		channelID,
		vs.SessionID,
	)
}

func OnVoiceServerUpdate(ev *discordgo.VoiceServerUpdate) {
	if manager == nil {
		return
	}
	endpoint := ""
	if ev.Endpoint != "" {
		endpoint = ev.Endpoint
	}

	manager.client.OnVoiceServerUpdate(
		context.Background(),
		snowflake.MustParse(ev.GuildID),
		ev.Token,
		endpoint,
	)
}

func getOrCreateGuildPlayer(guildID string) *guildPlayer {
	manager.mu.Lock()
	defer manager.mu.Unlock()

	if gp, ok := manager.players[guildID]; ok {
		return gp
	}

	player := manager.client.Player(snowflake.MustParse(guildID))
	gp := &guildPlayer{
		Player: player,
		Queue:  make([]TrackEntry, 0),
		Volume: 100,
	}
	manager.players[guildID] = gp
	return gp
}

func EnqueueAndPlay(guildID, voiceChannelID, textChannelID, query, requester string) (*TrackEntry, bool, error) {
	if manager == nil {
		return nil, false, fmt.Errorf("music manager not initialized")
	}

	// Join voice channel
	if err := manager.session.ChannelVoiceJoinManual(guildID, voiceChannelID, false, false); err != nil {
		return nil, false, fmt.Errorf("join voice channel: %w", err)
	}

	gp := getOrCreateGuildPlayer(guildID)
	if textChannelID != "" {
		manager.mu.Lock()
		gp.TextChannelID = textChannelID
		manager.mu.Unlock()
	}

	var loadedTracks []lavalink.Track

	manager.client.BestNode().LoadTracksHandler(
		context.Background(),
		query,
		disgolink.NewResultHandler(
			func(track lavalink.Track) {
				loadedTracks = append(loadedTracks, track)
			},
			func(playlist lavalink.Playlist) {
				if len(playlist.Tracks) == 0 {
					return
				}

				// Respect Lavalink's selectedTrack when present by rotating the playlist.
				selected := playlist.Info.SelectedTrack
				if selected >= 0 && selected < len(playlist.Tracks) {
					loadedTracks = append(loadedTracks, playlist.Tracks[selected:]...)
					loadedTracks = append(loadedTracks, playlist.Tracks[:selected]...)
					return
				}

				loadedTracks = append(loadedTracks, playlist.Tracks...)
			},
			func(tracks []lavalink.Track) {
				if len(tracks) > 0 {
					loadedTracks = append(loadedTracks, tracks[0])
				}
			},
			func() {
			},
			func(err error) {
			},
		),
	)

	if len(loadedTracks) == 0 {
		return nil, false, fmt.Errorf("no tracks found for query")
	}

	// Some sources may return unplayable entries; pick the first with an encoded track.
	firstPlayableIdx := 0
	for idx := range loadedTracks {
		if loadedTracks[idx].Encoded != "" {
			firstPlayableIdx = idx
			break
		}
	}
	if firstPlayableIdx != 0 && firstPlayableIdx < len(loadedTracks) {
		loadedTracks = append(loadedTracks[firstPlayableIdx:], loadedTracks[:firstPlayableIdx]...)
	}

	entries := make([]TrackEntry, len(loadedTracks))
	for idx, t := range loadedTracks {
		entries[idx] = TrackEntry{
			Track:       t,
			RequestedBy: requester,
		}
	}
	firstEntry := entries[0]

	manager.mu.Lock()
	shouldStart := len(gp.Queue) == 0
	gp.Queue = append(gp.Queue, entries...)
	gp.IdleSince = time.Time{}
	manager.mu.Unlock()

	if shouldStart {
		if err := gp.Player.Update(context.Background(), lavalink.WithTrack(firstEntry.Track)); err != nil {
			return nil, false, fmt.Errorf("start track: %w", err)
		}
		if gp.TextChannelID != "" {
			postNowPlaying(guildID, gp.TextChannelID, firstEntry)
		}
	}

	return &firstEntry, shouldStart, nil
}

func onTrackEnd(player disgolink.Player, event lavalink.TrackEndEvent) {
	if manager == nil {
		return
	}
	if !event.Reason.MayStartNext() {
		return
	}

	guildID := player.GuildID().String()

	manager.mu.Lock()
	gp, ok := manager.players[guildID]
	if !ok || len(gp.Queue) == 0 {
		manager.mu.Unlock()
		return
	}

	// repeat current track
	if gp.RepeatSong {
		next := gp.Queue[0]
		textChannelID := gp.TextChannelID
		manager.mu.Unlock()

		if err := gp.Player.Update(context.Background(), lavalink.WithTrack(next.Track)); err != nil {
			return
		}
		if textChannelID != "" {
			postNowPlaying(guildID, textChannelID, next)
		}
		return
	}

	// pop current and optionally cycle it to end
	finished := gp.Queue[0]
	gp.Queue = gp.Queue[1:]
	if gp.RepeatQueue {
		gp.Queue = append(gp.Queue, finished)
	}

	if len(gp.Queue) == 0 {
		textChannelID := gp.TextChannelID
		nowPlayingMsgID := gp.NowPlayingMsgID
		gp.NowPlayingMsgID = ""
		gp.IdleSince = time.Now()
		manager.mu.Unlock()
		if textChannelID != "" && nowPlayingMsgID != "" {
			disableNowPlaying(textChannelID, nowPlayingMsgID)
		}
		_ = gp.Player.Update(context.Background(), lavalink.WithNullTrack())
		return
	}

	next := gp.Queue[0]
	textChannelID := gp.TextChannelID
	manager.mu.Unlock()

	if err := gp.Player.Update(context.Background(), lavalink.WithTrack(next.Track)); err != nil {
		return
	}
	if textChannelID != "" {
		postNowPlaying(guildID, textChannelID, next)
	}
}

func Pause(guildID string, pause bool) error {
	if manager == nil {
		return fmt.Errorf("music manager not initialized")
	}
	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	gp.Paused = pause
	if pause {
		gp.IdleSince = time.Time{}
	}
	manager.mu.Unlock()
	return gp.Player.Update(context.Background(), lavalink.WithPaused(pause))
}

func Skip(guildID string) error {
	if manager == nil {
		return fmt.Errorf("music manager not initialized")
	}
	gp := getOrCreateGuildPlayer(guildID)

	manager.mu.Lock()
	if len(gp.Queue) == 0 {
		manager.mu.Unlock()
		return nil
	}
	// skip always advances regardless of repeat song
	skipped := gp.Queue[0]
	gp.Queue = gp.Queue[1:]
	if gp.RepeatQueue {
		gp.Queue = append(gp.Queue, skipped)
	}

	if len(gp.Queue) == 0 {
		textChannelID := gp.TextChannelID
		nowPlayingMsgID := gp.NowPlayingMsgID
		gp.NowPlayingMsgID = ""
		gp.IdleSince = time.Now()
		manager.mu.Unlock()
		if textChannelID != "" && nowPlayingMsgID != "" {
			disableNowPlaying(textChannelID, nowPlayingMsgID)
		}
		return gp.Player.Update(context.Background(), lavalink.WithNullTrack())
	}

	next := gp.Queue[0]
	textChannelID := gp.TextChannelID
	manager.mu.Unlock()

	if err := gp.Player.Update(context.Background(), lavalink.WithTrack(next.Track)); err != nil {
		return err
	}
	if textChannelID != "" {
		postNowPlaying(guildID, textChannelID, next)
	}
	return nil
}

func Stop(guildID string) error {
	if manager == nil {
		return fmt.Errorf("music manager not initialized")
	}
	gp := getOrCreateGuildPlayer(guildID)

	manager.mu.Lock()
	textChannelID := gp.TextChannelID
	nowPlayingMsgID := gp.NowPlayingMsgID
	gp.Queue = nil
	gp.Paused = false
	gp.NowPlayingMsgID = ""
	gp.IdleSince = time.Now()
	manager.mu.Unlock()

	if textChannelID != "" && nowPlayingMsgID != "" {
		disableNowPlaying(textChannelID, nowPlayingMsgID)
	}

	return gp.Player.Update(context.Background(), lavalink.WithNullTrack())
}

func GetQueue(guildID string) (*TrackEntry, []TrackEntry) {
	if manager == nil {
		return nil, nil
	}
	gp := getOrCreateGuildPlayer(guildID)

	manager.mu.Lock()
	defer manager.mu.Unlock()

	if len(gp.Queue) == 0 {
		return nil, nil
	}

	current := gp.Queue[0]
	rest := make([]TrackEntry, len(gp.Queue)-1)
	copy(rest, gp.Queue[1:])
	return &current, rest
}

func ToggleRepeatSong(guildID string) (bool, bool) {
	if manager == nil {
		return false, false
	}
	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	defer manager.mu.Unlock()
	gp.RepeatSong = !gp.RepeatSong
	if gp.RepeatSong {
		gp.RepeatQueue = false
	}
	return gp.RepeatSong, gp.RepeatQueue
}

func ToggleRepeatQueue(guildID string) (bool, bool) {
	if manager == nil {
		return false, false
	}
	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	defer manager.mu.Unlock()
	gp.RepeatQueue = !gp.RepeatQueue
	if gp.RepeatQueue {
		gp.RepeatSong = false
	}
	return gp.RepeatSong, gp.RepeatQueue
}

func repeatFlags(guildID string) (bool, bool) {
	if manager == nil {
		return false, false
	}
	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	defer manager.mu.Unlock()
	return gp.RepeatSong, gp.RepeatQueue
}

func pausedAndVolume(guildID string) (bool, int) {
	if manager == nil {
		return false, 100
	}
	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	defer manager.mu.Unlock()
	vol := gp.Volume
	if vol <= 0 {
		vol = 100
	}
	return gp.Paused, vol
}

func SetVolume(guildID string, volume int) error {
	if manager == nil {
		return fmt.Errorf("music manager not initialized")
	}
	if volume < 0 {
		volume = 0
	}
	if volume > 150 {
		volume = 150
	}

	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	gp.Volume = volume
	textChannelID := gp.TextChannelID
	nowPlayingMsgID := gp.NowPlayingMsgID
	manager.mu.Unlock()

	if err := gp.Player.Update(context.Background(), lavalink.WithVolume(volume)); err != nil {
		return err
	}

	// Refresh the embed so the displayed volume matches.
	if textChannelID != "" && nowPlayingMsgID != "" {
		current, _ := GetQueue(guildID)
		if current != nil {
			postNowPlaying(guildID, textChannelID, *current)
		}
	}
	return nil
}

func CurrentNowPlayingMessageID(guildID string) string {
	if manager == nil {
		return ""
	}
	gp := getOrCreateGuildPlayer(guildID)
	manager.mu.Lock()
	defer manager.mu.Unlock()
	return gp.NowPlayingMsgID
}