import Foundation

enum VideoShuffle {
    /// Groups items by key, shuffling within each group.
    ///
    /// Items whose key is `nil` are grouped under the empty string `""` so they
    /// form their own bucket rather than being dropped.
    static func groupByModel<T>(
        _ items: [T],
        keyFor: (T) -> String?,
        using rng: inout some RandomNumberGenerator
    ) -> [String: [T]] {
        var groups: [String: [T]] = [:]
        for item in items {
            let key = keyFor(item) ?? ""
            groups[key, default: []].append(item)
        }
        for key in groups.keys {
            groups[key]!.shuffle(using: &rng)
        }
        return groups
    }

    /// Picks a random item from a random non-empty bucket, avoiding
    /// `excludingKey` when possible.
    ///
    /// Returned buckets are a mutated copy of the input with the picked item
    /// removed (pure — input is not modified). When the excluded bucket is the
    /// only non-empty one, it is still picked (falling back gracefully).
    ///
    /// - Returns: the picked item, the key of the bucket it came from, and the
    ///   updated bucket dictionary — or `nil` if all buckets are empty.
    static func pickRandomFromBuckets<T>(
        _ buckets: [String: [T]],
        excludingKey: String? = nil,
        using rng: inout some RandomNumberGenerator
    ) -> (item: T, key: String, remaining: [String: [T]])? {
        // Collect keys with at least one item. Sort for deterministic ordering
        // before the random pick so test output is reproducible when a seeded
        // RNG is used.
        let nonEmptyKeys = buckets
            .filter { !$0.value.isEmpty }
            .keys
            .sorted()
        guard !nonEmptyKeys.isEmpty else { return nil }

        // Prefer keys other than the excluded one, but fall back if excluding
        // would leave no candidates.
        let candidateKeys: [String]
        if let excludingKey, nonEmptyKeys.contains(excludingKey), nonEmptyKeys.count > 1 {
            candidateKeys = nonEmptyKeys.filter { $0 != excludingKey }
        } else {
            candidateKeys = nonEmptyKeys
        }

        guard let pickedKey = candidateKeys.randomElement(using: &rng),
              var items = buckets[pickedKey],
              !items.isEmpty else {
            return nil
        }

        // Items were shuffled during groupByModel, so removeFirst is a
        // uniform random pick among the remaining items in this bucket.
        let item = items.removeFirst()
        var remaining = buckets
        remaining[pickedKey] = items
        return (item, pickedKey, remaining)
    }

    /// Convenience overload using the system random generator.
    static func pickRandomFromBuckets<T>(
        _ buckets: [String: [T]],
        excludingKey: String? = nil
    ) -> (item: T, key: String, remaining: [String: [T]])? {
        var rng = SystemRandomNumberGenerator()
        return pickRandomFromBuckets(buckets, excludingKey: excludingKey, using: &rng)
    }
}