import Foundation

public enum WindowSpaceMappingPolicy {
    public static func mergedSpaceIDsByWindowID(
        requestedWindowIDs: [UInt32],
        directSpaceIDsByWindowID: [UInt32: [UInt64]],
        windowIDsBySpaceID: [UInt64: [UInt32]],
        orderedSpaceIDs: [UInt64]
    ) -> [UInt32: [UInt64]] {
        let requestedWindowIDSet = Set(requestedWindowIDs)
        var result: [UInt32: [UInt64]] = [:]

        for windowID in requestedWindowIDs {
            let directSpaceIDs = orderedUnique(directSpaceIDsByWindowID[windowID] ?? [])
            if !directSpaceIDs.isEmpty {
                result[windowID] = directSpaceIDs
            }
        }

        for spaceID in normalizedSpaceOrder(
            orderedSpaceIDs: orderedSpaceIDs,
            windowIDsBySpaceID: windowIDsBySpaceID
        ) {
            for windowID in windowIDsBySpaceID[spaceID] ?? [] {
                guard requestedWindowIDSet.contains(windowID),
                      result[windowID] == nil
                else {
                    continue
                }

                result[windowID, default: []].append(spaceID)
            }
        }

        return result.mapValues(orderedUnique)
    }

    private static func normalizedSpaceOrder(
        orderedSpaceIDs: [UInt64],
        windowIDsBySpaceID: [UInt64: [UInt32]]
    ) -> [UInt64] {
        let orderedUniqueSpaceIDs = orderedUnique(orderedSpaceIDs)
        if !orderedUniqueSpaceIDs.isEmpty {
            return orderedUniqueSpaceIDs
        }

        return windowIDsBySpaceID.keys.sorted()
    }

    private static func orderedUnique<T: Hashable>(_ values: [T]) -> [T] {
        var seen = Set<T>()
        return values.filter { seen.insert($0).inserted }
    }
}