ersatztv/ErsatzTV.Core.Tests/Scheduling/WeightedShuffleCollectionEnumeratorTests.cs

using ErsatzTV.Core.Domain;
using ErsatzTV.Core.Scheduling;
using LanguageExt.UnsafeValueAccess;
using NUnit.Framework;
using Shouldly;

namespace ErsatzTV.Core.Tests.Scheduling;

/// <summary>
///     Pins the weighted / fair-share distribution contract (#70). The sequence semantics here are the
///     product decision, so they are asserted exactly rather than statistically.
/// </summary>
[TestFixture]
public class WeightedShuffleCollectionEnumeratorTests
{
    // ids are allocated per source so an emitted item's source is identifiable from its id
    private const int SourceAFirstId = 100;
    private const int SourceBFirstId = 200;
    private const int SourceCFirstId = 300;

    private static CollectionWithItems Source(string key, int firstId, int itemCount, int weight) =>
        new(
            0,
            0,
            key,
            Enumerable.Range(firstId, itemCount)
                .Select(i => new Movie { Id = i, MovieMetadata = [] })
                .Cast<MediaItem>()
                .ToList(),
            true,
            PlaybackOrder.WeightedShuffle,
            false,
            weight);

    private static string SourceOf(int id) => id switch
    {
        >= SourceCFirstId => "C",
        >= SourceBFirstId => "B",
        _ => "A"
    };

    private static List<string> TakeSourceSequence(WeightedShuffleCollectionEnumerator enumerator, int count)
    {
        var result = new List<string>();
        for (var i = 0; i < count; i++)
        {
            enumerator.Current.IsSome.ShouldBeTrue();
            result.Add(SourceOf(enumerator.Current.ValueUnsafe().Id));
            enumerator.MoveNext(Option<DateTimeOffset>.None);
        }

        return result;
    }

    [Test]
    public void Weights_Three_To_One_Emit_A_A_B_A()
    {
        // the canonical smooth-WRR contract: 3:1 spreads B through the rotation (A A B A),
        // rather than draining A first (A A A B) the way PlaylistItem.Count does
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 3, 3),
            Source("B", SourceBFirstId, 1, 1)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 1234, Index = 0 },
            CancellationToken.None);

        TakeSourceSequence(enumerator, 4).ShouldBe(["A", "A", "B", "A"]);
    }

    [Test]
    public void Equal_Weights_Are_Fair_Share_Regardless_Of_Collection_Size()
    {
        // the heart of goal (2): a 20-item source must air as often as a 2-item source.
        // ShuffleInOrder cannot do this -- it plays every item once, so airtime tracks size.
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 20, 1),
            Source("B", SourceBFirstId, 2, 1)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 1234, Index = 0 },
            CancellationToken.None);

        List<string> sequence = TakeSourceSequence(enumerator, 40);

        // equal weights => strict alternation, so the small source loops rather than falling silent
        sequence.Count(s => s == "A").ShouldBe(20);
        sequence.Count(s => s == "B").ShouldBe(20);
    }

    [Test]
    public void Ties_Break_To_The_Earliest_Source_In_List_Order()
    {
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 4, 1),
            Source("B", SourceBFirstId, 4, 1),
            Source("C", SourceCFirstId, 4, 1)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 99, Index = 0 },
            CancellationToken.None);

        // all accumulators tie every round, so list order decides
        TakeSourceSequence(enumerator, 6).ShouldBe(["A", "B", "C", "A", "B", "C"]);
    }

    [Test]
    public void An_Unweighted_Collection_Is_Fair_Share()
    {
        // a source arriving without an explicit weight must rotate as fair-share, never fall out
        var collections = new List<CollectionWithItems>
        {
            new(0, 0, "A", [new Movie { Id = SourceAFirstId, MovieMetadata = [] }], true, PlaybackOrder.WeightedShuffle, false),
            new(0, 0, "B", [new Movie { Id = SourceBFirstId, MovieMetadata = [] }], true, PlaybackOrder.WeightedShuffle, false)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 7, Index = 0 },
            CancellationToken.None);

        TakeSourceSequence(enumerator, 4).ShouldBe(["A", "B", "A", "B"]);
    }

    [Test]
    [TestCase(0)]
    [TestCase(-5)]
    public void A_Non_Positive_Weight_Does_Not_Delete_The_Source(int weight)
    {
        // the write path bounds weight, but a row can predate that gate. Treating a 0/negative weight as a
        // filter would remove the source from the channel silently -- the exact failure this order avoids
        // everywhere else. It is clamped to the floor instead.
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 2, 1),
            Source("B", SourceBFirstId, 2, weight)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 21, Index = 0 },
            CancellationToken.None);

        TakeSourceSequence(enumerator, 8).ShouldContain("B");
    }

    [Test]
    public void An_Enormous_Weight_Does_Not_Overflow_The_Rotation()
    {
        // summing unclamped weights is checked arithmetic, so an out-of-range row would throw from inside a
        // playout build rather than merely schedule oddly
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 2, int.MaxValue),
            Source("B", SourceBFirstId, 2, int.MaxValue)
        };

        WeightedShuffleCollectionEnumerator enumerator = null;
        Should.NotThrow(() => enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 5, Index = 0 },
            CancellationToken.None));

        // both clamp to the ceiling, so they tie and alternate by list order
        TakeSourceSequence(enumerator, 4).ShouldBe(["A", "B", "A", "B"]);
    }

    [Test]
    [TestCase(12)]
    [TestCase(13)]
    [TestCase(20)]
    [TestCase(37)]
    public void Restoring_Past_A_Rotation_Wrap_Equals_Advancing_To_It(int target)
    {
        // the stateless claim has to hold ACROSS a wrap, not just within the first rotation: each wrap
        // re-derives the rotation from the new seed alone, so (Seed, Index) still determines position.
        // Cycle length here is 12, so every case but the first crosses at least one wrap.
        List<CollectionWithItems> Collections() =>
        [
            Source("A", SourceAFirstId, 5, 3),
            Source("B", SourceBFirstId, 3, 1)
        ];

        const int Seed = 909;

        var advanced = new WeightedShuffleCollectionEnumerator(
            Collections(),
            new CollectionEnumeratorState { Seed = Seed, Index = 0 },
            CancellationToken.None);
        for (var i = 0; i < target; i++)
        {
            advanced.MoveNext(Option<DateTimeOffset>.None);
        }

        var restored = new WeightedShuffleCollectionEnumerator(
            Collections(),
            new CollectionEnumeratorState { Seed = advanced.State.Seed, Index = advanced.State.Index },
            CancellationToken.None);

        restored.Current.ValueUnsafe().Id.ShouldBe(advanced.Current.ValueUnsafe().Id);
    }

    [Test]
    public void Restoring_At_An_Index_Equals_Advancing_To_It()
    {
        // the stateless contract: (Seed, Index) fully determines position, which is what lets the
        // existing CollectionEnumeratorState persistence carry this order with no per-source counters
        List<CollectionWithItems> Collections() =>
        [
            Source("A", SourceAFirstId, 5, 3),
            Source("B", SourceBFirstId, 3, 1)
        ];

        const int Seed = 4242;
        const int Target = 7;

        var advanced = new WeightedShuffleCollectionEnumerator(
            Collections(),
            new CollectionEnumeratorState { Seed = Seed, Index = 0 },
            CancellationToken.None);
        for (var i = 0; i < Target; i++)
        {
            advanced.MoveNext(Option<DateTimeOffset>.None);
        }

        var restored = new WeightedShuffleCollectionEnumerator(
            Collections(),
            new CollectionEnumeratorState { Seed = Seed, Index = Target },
            CancellationToken.None);

        restored.State.Index.ShouldBe(advanced.State.Index);
        restored.State.Seed.ShouldBe(advanced.State.Seed);
        restored.Current.ValueUnsafe().Id.ShouldBe(advanced.Current.ValueUnsafe().Id);
    }

    [Test]
    public void A_Single_Source_Emits_All_Of_Its_Items()
    {
        var collections = new List<CollectionWithItems> { Source("A", SourceAFirstId, 5, 3) };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 11, Index = 0 },
            CancellationToken.None);

        var seen = new System.Collections.Generic.HashSet<int>();
        for (var i = 0; i < 5; i++)
        {
            seen.Add(enumerator.Current.ValueUnsafe().Id);
            enumerator.MoveNext(Option<DateTimeOffset>.None);
        }

        seen.Count.ShouldBe(5);
    }

    [Test]
    public void An_Empty_Source_Is_Ignored_Rather_Than_Emitting_Nothing()
    {
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 2, 1),
            new(0, 0, "empty", [], true, PlaybackOrder.WeightedShuffle, false, 5)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 3, Index = 0 },
            CancellationToken.None);

        // a heavily-weighted empty source must not starve the rotation or emit None
        TakeSourceSequence(enumerator, 4).ShouldBe(["A", "A", "A", "A"]);
    }

    [Test]
    [Timeout(10_000)]
    public void Single_Item_Sources_Do_Not_Hang_On_Rotation_Wrap()
    {
        // regression: the wrap retries a rebuild to avoid an immediate repeat, but this order's lead item is
        // decided by weight, so the heaviest source always leads. With one item in it the lead is invariant and
        // an unbounded retry never terminates -- a hung playout build, not a wrong one. Found by the
        // non-vacuity control, which hung instead of failing.
        var collections = new List<CollectionWithItems>
        {
            Source("A", SourceAFirstId, 1, 3),
            Source("B", SourceBFirstId, 1, 1)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 1234, Index = 0 },
            CancellationToken.None);

        // walk several full rotations so the wrap path is exercised repeatedly
        List<string> sequence = TakeSourceSequence(enumerator, 24);

        sequence.ShouldContain("A");
        sequence.ShouldContain("B");
    }

    [Test]
    public void No_Sources_Yields_No_Current()
    {
        var enumerator = new WeightedShuffleCollectionEnumerator(
            [],
            new CollectionEnumeratorState { Seed = 1, Index = 0 },
            CancellationToken.None);

        enumerator.Current.IsNone.ShouldBeTrue();
        enumerator.Count.ShouldBe(0);
        Should.NotThrow(() => enumerator.MoveNext(Option<DateTimeOffset>.None));
    }

    [Test]
    public void A_Custom_Ordered_Source_Keeps_Its_Order()
    {
        var collections = new List<CollectionWithItems>
        {
            new(
                0,
                0,
                "A",
                Enumerable.Range(SourceAFirstId, 4)
                    .Select(i => new Movie { Id = i, MovieMetadata = [] })
                    .Cast<MediaItem>()
                    .ToList(),
                true,
                PlaybackOrder.WeightedShuffle,
                true,
                1)
        };

        var enumerator = new WeightedShuffleCollectionEnumerator(
            collections,
            new CollectionEnumeratorState { Seed = 555, Index = 0 },
            CancellationToken.None);

        var ids = new List<int>();
        for (var i = 0; i < 4; i++)
        {
            ids.Add(enumerator.Current.ValueUnsafe().Id);
            enumerator.MoveNext(Option<DateTimeOffset>.None);
        }

        ids.ShouldBe([SourceAFirstId, SourceAFirstId + 1, SourceAFirstId + 2, SourceAFirstId + 3]);
    }
}