MLBApp/mlbTVOSTests/VideoShuffleTests.swift

import Foundation
import Testing
@testable import mlbTVOS

struct VideoShuffleTests {
    // MARK: - Helpers

    private struct Clip: Equatable, Hashable {
        let model: String
        let id: Int
    }

    private static func makeClips(_ counts: [String: Int]) -> [Clip] {
        var result: [Clip] = []
        var idCounter = 0
        for (model, count) in counts.sorted(by: { $0.key < $1.key }) {
            for _ in 0..<count {
                result.append(Clip(model: model, id: idCounter))
                idCounter += 1
            }
        }
        return result
    }

    /// Returns the length of the longest contiguous run of same-model clips.
    private static func maxRunLength(_ clips: [Clip]) -> Int {
        guard !clips.isEmpty else { return 0 }
        var maxRun = 1
        var currentRun = 1
        for i in 1..<clips.count {
            if clips[i].model == clips[i - 1].model {
                currentRun += 1
                maxRun = max(maxRun, currentRun)
            } else {
                currentRun = 1
            }
        }
        return maxRun
    }

    /// A deterministic RNG so tests that depend on randomness are reproducible.
    private struct SeededRNG: RandomNumberGenerator {
        private var state: UInt64
        init(seed: UInt64) { self.state = seed == 0 ? 0xdeadbeef : seed }
        mutating func next() -> UInt64 {
            state ^= state << 13
            state ^= state >> 7
            state ^= state << 17
            return state
        }
    }

    // MARK: - groupByModel

    @Test func groupByModelBucketsCorrectly() {
        let clips = Self.makeClips(["a": 3, "b": 2, "c": 1])
        var rng = SeededRNG(seed: 1)
        let groups = VideoShuffle.groupByModel(
            clips,
            keyFor: { $0.model },
            using: &rng
        )
        #expect(groups["a"]?.count == 3)
        #expect(groups["b"]?.count == 2)
        #expect(groups["c"]?.count == 1)
        #expect(groups.count == 3)
    }

    @Test func groupByModelHandlesEmptyInput() {
        var rng = SeededRNG(seed: 1)
        let groups = VideoShuffle.groupByModel(
            [Clip](),
            keyFor: { $0.model },
            using: &rng
        )
        #expect(groups.isEmpty)
    }

    @Test func groupByModelNilKeysGoToEmptyStringBucket() {
        struct Item { let key: String?; let id: Int }
        let items = [
            Item(key: "a", id: 0),
            Item(key: nil, id: 1),
            Item(key: nil, id: 2),
        ]
        var rng = SeededRNG(seed: 1)
        let groups = VideoShuffle.groupByModel(
            items,
            keyFor: { $0.key },
            using: &rng
        )
        #expect(groups["a"]?.count == 1)
        #expect(groups[""]?.count == 2)
    }

    // MARK: - pickRandomFromBuckets

    @Test func pickRandomFromBucketsReturnsNilWhenAllEmpty() {
        let buckets: [String: [Clip]] = ["a": [], "b": []]
        var rng = SeededRNG(seed: 1)
        let result = VideoShuffle.pickRandomFromBuckets(
            buckets,
            using: &rng
        )
        #expect(result == nil)
    }

    @Test func pickRandomFromBucketsRemovesPickedItem() {
        let buckets: [String: [Clip]] = [
            "a": [Clip(model: "a", id: 0), Clip(model: "a", id: 1)],
            "b": [Clip(model: "b", id: 2)],
        ]
        var rng = SeededRNG(seed: 1)
        let result = VideoShuffle.pickRandomFromBuckets(
            buckets,
            using: &rng
        )
        #expect(result != nil)
        guard let pick = result else { return }
        let originalCount = buckets.values.reduce(0) { $0 + $1.count }
        let remainingCount = pick.remaining.values.reduce(0) { $0 + $1.count }
        #expect(remainingCount == originalCount - 1)
    }

    @Test func pickRandomFromBucketsRespectsExclusion() {
        // With two non-empty buckets, excluding "a" should always return "b".
        let buckets: [String: [Clip]] = [
            "a": [Clip(model: "a", id: 0), Clip(model: "a", id: 1)],
            "b": [Clip(model: "b", id: 2), Clip(model: "b", id: 3)],
        ]
        var rng = SeededRNG(seed: 1)
        for _ in 0..<20 {
            let result = VideoShuffle.pickRandomFromBuckets(
                buckets,
                excludingKey: "a",
                using: &rng
            )
            #expect(result?.key == "b")
        }
    }

    @Test func pickRandomFromBucketsFallsBackWhenOnlyExcludedIsNonEmpty() {
        let buckets: [String: [Clip]] = [
            "a": [Clip(model: "a", id: 0)],
            "b": [],
        ]
        var rng = SeededRNG(seed: 1)
        let result = VideoShuffle.pickRandomFromBuckets(
            buckets,
            excludingKey: "a",
            using: &rng
        )
        #expect(result?.key == "a", "should fall back to 'a' since 'b' is empty")
    }

    // MARK: - End-to-end: never same model back-to-back

    /// Simulates a full playback session on the real NSFW distribution and
    /// verifies that no two consecutive picks share a model. This is the
    /// invariant the user actually cares about.
    @Test func realNSFWDistributionNeverBackToBack() {
        let clips = Self.makeClips([
            "kayla-lauren": 502,
            "tabatachang": 147,
            "werkout": 137,
            "ray-mattos": 108,
            "dani-speegle-2": 49,
            "josie-hamming-2": 13,
        ])
        #expect(clips.count == 956)

        var rng = SeededRNG(seed: 42)
        var buckets = VideoShuffle.groupByModel(
            clips,
            keyFor: { $0.model },
            using: &rng
        )

        // Simulate 2000 picks — more than the total clip count so we exercise
        // the refill logic for the small buckets (josie-hamming-2 has only 13).
        var sequence: [Clip] = []
        var lastModel: String? = nil
        for _ in 0..<2000 {
            // Refill any empty buckets by reshuffling the original clips for
            // that model (mirrors the production logic in GamesViewModel).
            for key in Array(buckets.keys) where buckets[key]?.isEmpty ?? true {
                let modelClips = clips.filter { $0.model == key }
                buckets[key] = modelClips.shuffled(using: &rng)
            }

            guard let pick = VideoShuffle.pickRandomFromBuckets(
                buckets,
                excludingKey: lastModel,
                using: &rng
            ) else {
                Issue.record("pick returned nil")
                return
            }
            buckets = pick.remaining
            sequence.append(pick.item)
            lastModel = pick.key
        }

        let run = Self.maxRunLength(sequence)
        #expect(run == 1, "expected max run 1, got \(run)")
    }

    @Test func modelDistributionIsRoughlyUniform() {
        // Over a long session, each model should be picked roughly
        // (totalPicks / modelCount) times — not weighted by bucket size.
        let clips = Self.makeClips([
            "kayla-lauren": 502,
            "tabatachang": 147,
            "werkout": 137,
            "ray-mattos": 108,
            "dani-speegle-2": 49,
            "josie-hamming-2": 13,
        ])

        var rng = SeededRNG(seed: 99)
        var buckets = VideoShuffle.groupByModel(
            clips,
            keyFor: { $0.model },
            using: &rng
        )

        let totalPicks = 6000
        var modelCounts: [String: Int] = [:]
        var lastModel: String? = nil
        for _ in 0..<totalPicks {
            for key in Array(buckets.keys) where buckets[key]?.isEmpty ?? true {
                let modelClips = clips.filter { $0.model == key }
                buckets[key] = modelClips.shuffled(using: &rng)
            }
            guard let pick = VideoShuffle.pickRandomFromBuckets(
                buckets,
                excludingKey: lastModel,
                using: &rng
            ) else { continue }
            buckets = pick.remaining
            modelCounts[pick.key, default: 0] += 1
            lastModel = pick.key
        }

        // With 6 models and 6000 picks, expected per model ≈ 1000.
        // Allow ±25% (750 to 1250) since the exclusion rule introduces
        // slight bias against whichever model was just picked.
        for (model, count) in modelCounts {
            #expect(count >= 750 && count <= 1250, "\(model) got \(count) picks")
        }
    }
}