minstrel/internal/playlists/discover.go

package playlists

import (
	"context"
	"fmt"

	"github.com/jackc/pgx/v5/pgtype"

	"git.fabledsword.com/bvandeusen/minstrel/internal/db/dbq"
)

const (
	discoverTotalSlots         = 100
	discoverDormantSlots       = 40
	discoverCrossUserSlots     = 30
	discoverRandomSlots        = 30
	discoverMaxTracksPerAlbum  = 2
	discoverMaxTracksPerArtist = 3
)

// discoverTrack is the common shape used by the bucket allocator. The
// three sqlc-generated row types collapse into this internal struct so
// downstream functions don't need to be generic over them.
type discoverTrack struct {
	ID       pgtype.UUID
	AlbumID  pgtype.UUID
	ArtistID pgtype.UUID
}

// buildDiscoverCandidates assembles the Discover playlist track list.
// Pulls from three buckets, applies per-album/per-artist caps, then
// redistributes any deficit equally across the remaining buckets.
//
// Returns up to discoverTotalSlots track IDs in the order they should
// appear in the playlist (round-robin interleaved across buckets).
// Empty slice when the library has no eligible tracks at all.
func buildDiscoverCandidates(ctx context.Context, q *dbq.Queries, userID pgtype.UUID, dateStr string) ([]rankedCandidate, error) {
	dormantRows, err := q.ListDormantArtistTracksForDiscover(ctx, dbq.ListDormantArtistTracksForDiscoverParams{
		UserID: userID, Column2: dateStr,
	})
	if err != nil {
		return nil, fmt.Errorf("dormant bucket: %w", err)
	}
	crossUserRows, err := q.ListCrossUserLikedTracksForDiscover(ctx, dbq.ListCrossUserLikedTracksForDiscoverParams{
		UserID: userID, Column2: dateStr,
	})
	if err != nil {
		return nil, fmt.Errorf("cross-user bucket: %w", err)
	}
	randomRows, err := q.ListRandomUnheardTracksForDiscover(ctx, dbq.ListRandomUnheardTracksForDiscoverParams{
		UserID: userID, Column2: dateStr,
	})
	if err != nil {
		return nil, fmt.Errorf("random bucket: %w", err)
	}

	// Adapt sqlc-generated rows into the internal struct, then apply
	// per-album / per-artist caps.
	dormantPool := capByAlbumAndArtist(dormantRowsToTracks(dormantRows))
	crossUserPool := capByAlbumAndArtist(crossUserRowsToTracks(crossUserRows))
	randomPool := capByAlbumAndArtist(randomRowsToTracks(randomRows))

	// Allocate slots with redistribution.
	allocations := redistributeSlots([]bucketRequest{
		{want: discoverDormantSlots, available: len(dormantPool)},
		{want: discoverCrossUserSlots, available: len(crossUserPool)},
		{want: discoverRandomSlots, available: len(randomPool)},
	})

	// Take the head of each bucket's capped pool.
	dormantTake := dormantPool[:allocations[0]]
	crossUserTake := crossUserPool[:allocations[1]]
	randomTake := randomPool[:allocations[2]]

	// Round-robin interleave so the playlist doesn't front-load one
	// flavour.
	out := interleaveBuckets(dormantTake, crossUserTake, randomTake)

	// Convert to rankedCandidate (the type insertSystemPlaylist accepts).
	// Score is unused for Discover — the daily-deterministic md5 ordering
	// already gave us the ranking.
	ranked := make([]rankedCandidate, 0, len(out))
	for _, t := range out {
		ranked = append(ranked, rankedCandidate{TrackID: t.ID})
	}
	return ranked, nil
}

func dormantRowsToTracks(rows []dbq.ListDormantArtistTracksForDiscoverRow) []discoverTrack {
	out := make([]discoverTrack, len(rows))
	for i, r := range rows {
		out[i] = discoverTrack{ID: r.ID, AlbumID: r.AlbumID, ArtistID: r.ArtistID}
	}
	return out
}

func crossUserRowsToTracks(rows []dbq.ListCrossUserLikedTracksForDiscoverRow) []discoverTrack {
	out := make([]discoverTrack, len(rows))
	for i, r := range rows {
		out[i] = discoverTrack{ID: r.ID, AlbumID: r.AlbumID, ArtistID: r.ArtistID}
	}
	return out
}

func randomRowsToTracks(rows []dbq.ListRandomUnheardTracksForDiscoverRow) []discoverTrack {
	out := make([]discoverTrack, len(rows))
	for i, r := range rows {
		out[i] = discoverTrack{ID: r.ID, AlbumID: r.AlbumID, ArtistID: r.ArtistID}
	}
	return out
}

// capByAlbumAndArtist trims a bucket's pool so no single album appears
// more than discoverMaxTracksPerAlbum times and no single artist
// appears more than discoverMaxTracksPerArtist times. Preserves input
// order (which is the daily-deterministic md5 sort).
//
// pgtype.UUID's Bytes field is [16]byte, which IS comparable as a map
// key.
func capByAlbumAndArtist(rows []discoverTrack) []discoverTrack {
	albumCount := map[pgtype.UUID]int{}
	artistCount := map[pgtype.UUID]int{}
	out := make([]discoverTrack, 0, len(rows))
	for _, r := range rows {
		if albumCount[r.AlbumID] >= discoverMaxTracksPerAlbum {
			continue
		}
		if artistCount[r.ArtistID] >= discoverMaxTracksPerArtist {
			continue
		}
		albumCount[r.AlbumID]++
		artistCount[r.ArtistID]++
		out = append(out, r)
	}
	return out
}

// bucketRequest captures one bucket's desired and available count.
type bucketRequest struct {
	want      int
	available int
}

// redistributeSlots returns the final per-bucket allocation given each
// bucket's want and available count. When a bucket can't fill its want,
// the deficit splits equally between the OTHER buckets that still have
// supply. Sum of returned allocations <= sum of `want`.
//
// Algorithm:
//  1. Initial pass: allocate min(want, available) to each bucket;
//     compute supply = available - allocated.
//  2. For each deficit bucket (allocated < want), split the deficit
//     equally across peers that still have supply > 0. Track how
//     much of that deficit each pass actually placed (some peers may
//     not have enough supply for their full share).
//  3. Repeat until no movement happens — handles the case where a
//     peer's supply ran out mid-distribution and the residual needs
//     to roll to other peers in a later pass.
//
// Bug history (May 2026): an earlier version recomputed
// `deficit = b.want - final[i]` every pass without subtracting
// previously-redistributed amounts, so a bucket with permanently
// 0 final and want=30 kept handing out 30 to peers each pass. The
// `redistributed` array below tracks per-source absorbed deficit so
// subsequent passes only move the residual.
func redistributeSlots(buckets []bucketRequest) []int {
	n := len(buckets)
	final := make([]int, n)
	supply := make([]int, n)
	redistributed := make([]int, n)
	for i, b := range buckets {
		final[i] = b.want
		if final[i] > b.available {
			final[i] = b.available
		}
		supply[i] = b.available - final[i]
	}

	for pass := 0; pass < 4; pass++ {
		anyMoved := false
		for i, b := range buckets {
			// Remaining deficit = original gap minus what we've already
			// pushed out to peers in earlier passes.
			deficit := b.want - final[i] - redistributed[i]
			if deficit <= 0 {
				continue
			}
			peers := make([]int, 0, n-1)
			for j := 0; j < n; j++ {
				if j == i {
					continue
				}
				if supply[j] > 0 {
					peers = append(peers, j)
				}
			}
			if len(peers) == 0 {
				continue
			}
			share := deficit / len(peers)
			rem := deficit % len(peers)
			for k, j := range peers {
				take := share
				if k < rem {
					take++
				}
				if take > supply[j] {
					take = supply[j]
				}
				if take > 0 {
					final[j] += take
					supply[j] -= take
					redistributed[i] += take
					anyMoved = true
				}
			}
		}
		if !anyMoved {
			break
		}
	}

	// Final clamp: total <= discoverTotalSlots.
	sum := 0
	for _, f := range final {
		sum += f
	}
	for sum > discoverTotalSlots {
		// Trim the largest bucket. Should be rare.
		maxVal, idx := -1, -1
		for i, f := range final {
			if f > maxVal {
				maxVal, idx = f, i
			}
		}
		if idx < 0 {
			break
		}
		final[idx]--
		sum--
	}
	return final
}

// interleaveBuckets round-robins items from the buckets in order. The
// first item of each bucket appears before the second of any bucket.
func interleaveBuckets(buckets ...[]discoverTrack) []discoverTrack {
	total := 0
	for _, b := range buckets {
		total += len(b)
	}
	out := make([]discoverTrack, 0, total)
	indices := make([]int, len(buckets))
	for {
		anyAppended := false
		for bi, b := range buckets {
			if indices[bi] < len(b) {
				out = append(out, b[indices[bi]])
				indices[bi]++
				anyAppended = true
			}
		}
		if !anyAppended {
			break
		}
	}
	return out
}