ersatztv/ErsatzTV.Core/Scheduling/WeightedShuffleCollectionEnumerator.cs

using ErsatzTV.Core.Domain;
using ErsatzTV.Core.Extensions;
using ErsatzTV.Core.Interfaces.Scheduling;

namespace ErsatzTV.Core.Scheduling;

/// <summary>
///     Weighted / fair-share distribution (issue #70). Picks a *source* by smooth weighted round-robin,
///     then takes that source's next item. A source's <see cref="CollectionWithItems.Weight" /> is its
///     relative share of airtime: weights 3:1 emit <c>A A B A</c> per rotation.
///     <para>
///         Fair-share is the equal-weights degenerate case (the default, since Weight defaults to 1):
///         every source airs equally often regardless of library size, so a small show loops while a
///         large one works through its items. This is what distinguishes it from
///         <see cref="ShuffleInOrderCollectionEnumerator" />, whose balanced shuffle plays every item
///         exactly once per cycle and therefore leaves airtime proportional to collection size
///         (it prevents clumping, not domination).
///     </para>
///     <para>
///         Deterministic and stateless: the emitted sequence is a pure function of
///         (<see cref="CollectionEnumeratorState.Seed" />, <see cref="CollectionEnumeratorState.Index" />),
///         so it restores by replay exactly like its siblings and needs no per-source persisted counters.
///     </para>
/// </summary>
public class WeightedShuffleCollectionEnumerator : IMediaCollectionEnumerator
{
    /// <summary>
    ///     Upper bound on the precomputed rotation, so a large collection paired with a low weight
    ///     (picks needed grows as items * totalWeight / weight) can't allocate without limit. Hitting the
    ///     clamp only means the largest source doesn't finish its items within one rotation before the
    ///     reseed; the weight ratio itself is unaffected.
    /// </summary>
    private const int MaxCycleLength = 100_000;

    /// <summary>
    ///     Bounds the avoid-an-immediate-repeat retry when a rotation wraps. Unlike
    ///     <see cref="ShuffleInOrderCollectionEnumerator" />, whose reshuffle randomizes which item leads the
    ///     next cycle, this order's lead item is decided by weight — the heaviest source always wins the first
    ///     pick. So when that source has a single item the lead never changes and an unbounded retry would spin
    ///     forever. Avoiding a back-to-back repeat is a nicety; never terminating is not.
    /// </summary>
    private const int MaxReshuffleAttempts = 10;

    private readonly CancellationToken _cancellationToken;
    private readonly IList<CollectionWithItems> _collections;
    private readonly Lazy<Option<TimeSpan>> _lazyMinimumDuration;
    private Random _random;
    private MediaItem[] _shuffled;

    public WeightedShuffleCollectionEnumerator(
        IList<CollectionWithItems> collections,
        CollectionEnumeratorState state,
        CancellationToken cancellationToken)
    {
        CurrentIncludeInProgramGuide = Option<bool>.None;

        _collections = collections;
        _cancellationToken = cancellationToken;

        int cycleLength = CycleLength(collections);
        if (cycleLength > 0 && state.Index >= cycleLength)
        {
            state.Index = 0;
            state.Seed = new Random(state.Seed).Next();
        }

        _random = new Random(state.Seed);
        _shuffled = Build(_collections, _random);
        // computed over every source's items, not the current rotation: a rotation can be a strict subset
        // (see MaxCycleLength) and is rebuilt on every wrap, so caching over it would go stale
        _lazyMinimumDuration =
            new Lazy<Option<TimeSpan>>(() =>
                _collections
                    .Bind(c => c.MediaItems)
                    .Bind(i => i.GetNonZeroDuration())
                    .OrderBy(identity)
                    .HeadOrNone());

        State = new CollectionEnumeratorState { Seed = state.Seed };
        while (State.Index < state.Index)
        {
            MoveNext(Option<DateTimeOffset>.None);
        }
    }

    public void ResetState(CollectionEnumeratorState state)
    {
        // only rebuild if needed
        if (State.Seed != state.Seed)
        {
            _random = new Random(state.Seed);
            _shuffled = Build(_collections, _random);
        }

        State.Seed = state.Seed;
        State.Index = state.Index;
    }

    public string SchedulingContextName => "Weighted Shuffle";

    public CollectionEnumeratorState State { get; }

    public Option<MediaItem> Current => _shuffled.Length != 0 ? _shuffled[State.Index % _shuffled.Length] : None;

    public Option<bool> CurrentIncludeInProgramGuide { get; }

    public void MoveNext(Option<DateTimeOffset> scheduledAt)
    {
        if (_shuffled.Length == 0)
        {
            return;
        }

        if ((State.Index + 1) % _shuffled.Length == 0)
        {
            Option<MediaItem> tail = Current;

            State.Index = 0;
            var attempts = 0;
            do
            {
                State.Seed = _random.Next();
                _random = new Random(State.Seed);
                _shuffled = Build(_collections, _random);
                attempts++;
                // guard on the rotation, not the raw collection count: an empty source contributes nothing,
                // so a 2-collection/1-item rotation would otherwise burn every attempt chasing an
                // impossible non-repeat on every wrap
            } while (!_cancellationToken.IsCancellationRequested && _shuffled.Length > 1 &&
                     attempts < MaxReshuffleAttempts &&
                     Current.Map(x => x.Id) == tail.Map(x => x.Id));
        }
        else
        {
            State.Index++;
        }

        if (_shuffled.Length > 0)
        {
            State.Index %= _shuffled.Length;
        }
    }

    public Option<TimeSpan> MinimumDuration => _lazyMinimumDuration.Value;

    public int Count => _shuffled.Length;

    /// <summary>
    ///     Length of one rotation: enough picks for the source that needs the most of them to work through
    ///     all of its items at its share of the rotation. Sources needing fewer picks loop within the
    ///     rotation — that looping is exactly what makes equal weights mean equal airtime.
    /// </summary>
    private static int CycleLength(IList<CollectionWithItems> collections)
    {
        List<CollectionWithItems> active = ActiveSources(collections);
        if (active.Count == 0)
        {
            return 0;
        }

        long totalWeight = active.Sum(c => (long)EffectiveWeight(c));

        var length = 0;
        foreach (CollectionWithItems collection in active)
        {
            // picks needed for this source to emit every item once, given it wins Weight of every totalWeight picks
            var needed = (int)Math.Min(
                MaxCycleLength,
                Math.Ceiling(collection.MediaItems.Count * (double)totalWeight / EffectiveWeight(collection)));
            length = Math.Max(length, needed);
        }

        return length;
    }

    private static List<CollectionWithItems> ActiveSources(IList<CollectionWithItems> collections) =>
        collections.Filter(c => c.MediaItems.Count > 0).ToList();

    /// <summary>
    ///     Weight as the rotation should treat it. The write path bounds this (<see cref="MultiCollectionItemWeight" />),
    ///     but a row can predate that gate, so it is clamped rather than trusted: an out-of-range value must not
    ///     drop a source from the channel (0 or negative) or overflow the weight sum inside a playout build.
    /// </summary>
    private static int EffectiveWeight(CollectionWithItems collection) =>
        Math.Clamp(collection.Weight, MultiCollectionItemWeight.Minimum, MultiCollectionItemWeight.Maximum);

    private static MediaItem[] Build(IList<CollectionWithItems> collections, Random random)
    {
        List<CollectionWithItems> active = ActiveSources(collections);
        if (active.Count == 0)
        {
            return [];
        }

        var sources = active
            .Map(c => new WeightedSource { Weight = EffectiveWeight(c), Items = OrderItems(c, random) })
            .ToList();

        int totalWeight = sources.Sum(s => s.Weight);
        int cycleLength = CycleLength(collections);

        var result = new List<MediaItem>(cycleLength);
        for (var i = 0; i < cycleLength; i++)
        {
            // smooth weighted round-robin: every source gains its weight, the richest wins and pays the total.
            // Strict '>' keeps ties on the earliest source in list order, which makes the sequence deterministic.
            WeightedSource pick = null;
            foreach (WeightedSource source in sources)
            {
                source.Accumulator += source.Weight;
                if (pick is null || source.Accumulator > pick.Accumulator)
                {
                    pick = source;
                }
            }

            pick.Accumulator -= totalWeight;
            result.Add(pick.Items[pick.Cursor % pick.Items.Length]);
            pick.Cursor++;
        }

        return result.ToArray();
    }

    private static MediaItem[] OrderItems(CollectionWithItems collectionWithItems, Random random)
    {
        // A custom-ordered collection is an explicit user sequence, so it is honored rather than shuffled
        // (same rule as ShuffleInOrderCollectionEnumerator). Otherwise the source's items are shuffled per
        // rotation, which is what makes the reseed on wrap produce a different rotation.
        if (collectionWithItems.UseCustomOrder)
        {
            return collectionWithItems.MediaItems.ToArray();
        }

        return Shuffle(collectionWithItems.MediaItems, random);
    }

    private static MediaItem[] Shuffle(IEnumerable<MediaItem> list, Random random)
    {
        MediaItem[] copy = list.ToArray();

        int n = copy.Length;
        while (n > 1)
        {
            n--;
            int k = random.Next(n + 1);
            (copy[k], copy[n]) = (copy[n], copy[k]);
        }

        return copy;
    }

    private class WeightedSource
    {
        public int Weight { get; init; }
        public MediaItem[] Items { get; init; }
        public int Cursor { get; set; }
        public int Accumulator { get; set; }
    }
}