import Foundation

/// Decides which districts auto-expand into hub-and-spoke each tick.
///
/// Rules: a district expands only if it has at least two endpoints (a single-IP
/// org has nothing to fan out) and ranks among the busiest `topK` above
/// `floorBytes`. Once chosen it stays expanded for `ttlTicks` ticks, so a burst
/// leaves a smooth tail instead of flickering collapsed the moment it idles.
///
/// Pure and deterministic given its input — no SpriteKit, fully unit-testable.
public final class ExpansionPolicy {
    public struct District: Sendable {
        public let key: String
        public let bytes: UInt64
        public let endpointCount: Int
        public init(key: String, bytes: UInt64, endpointCount: Int) {
            self.key = key
            self.bytes = bytes
            self.endpointCount = endpointCount
        }
    }

    private let topK: Int
    private let ttlTicks: Int
    private let floorBytes: UInt64
    private var ttl: [String: Int] = [:]

    public init(topK: Int = 4, ttlTicks: Int = 4, floorBytes: UInt64 = 8 * 1024) {
        self.topK = topK
        self.ttlTicks = ttlTicks
        self.floorBytes = floorBytes
    }

    /// Advance one tick over the full set of currently-known districts (include
    /// idle ones so their tail can decay). Returns the set of keys to render
    /// expanded.
    public func update(_ districts: [District]) -> Set<String> {
        // Decay every existing timer.
        for key in ttl.keys { ttl[key] = max(0, (ttl[key] ?? 0) - 1) }

        // Refresh the busiest multi-endpoint districts.
        var winners = 0
        for d in districts.sorted(by: { $0.bytes > $1.bytes }) {
            guard winners < topK else { break }
            if d.bytes >= floorBytes && d.endpointCount >= 2 {
                ttl[d.key] = ttlTicks
                winners += 1
            }
        }

        // Expanded = live timer AND still has something to fan out.
        let endpointCounts = Dictionary(districts.map { ($0.key, $0.endpointCount) },
                                        uniquingKeysWith: { a, _ in a })
        var expanded = Set<String>()
        for (key, remaining) in ttl where remaining > 0 {
            if (endpointCounts[key] ?? 0) >= 2 { expanded.insert(key) }
        }
        return expanded
    }
}