import XCTest
@testable import NetworkCityCore

final class ExpansionPolicyTests: XCTestCase {
    typealias D = ExpansionPolicy.District

    func testBusyMultiEndpointOrgExpands() {
        let policy = ExpansionPolicy()
        let expanded = policy.update([D(key: "onedrive", bytes: 500_000, endpointCount: 15)])
        XCTAssertTrue(expanded.contains("onedrive"))
    }

    func testSingleEndpointNeverExpandsEvenIfHuge() {
        // The exact case we saw live: Google at 126 MB but only one active IP.
        let policy = ExpansionPolicy()
        let expanded = policy.update([D(key: "google", bytes: 126_000_000, endpointCount: 1)])
        XCTAssertTrue(expanded.isEmpty)
    }

    func testQuietOrgDoesNotExpand() {
        let policy = ExpansionPolicy(floorBytes: 8 * 1024)
        let expanded = policy.update([D(key: "apple", bytes: 100, endpointCount: 5)])
        XCTAssertTrue(expanded.isEmpty)
    }

    func testOnlyTopKExpand() {
        let policy = ExpansionPolicy(topK: 2)
        let districts = (0..<5).map { D(key: "org\($0)", bytes: UInt64(1000 - $0) * 1024, endpointCount: 3) }
        let expanded = policy.update(districts)
        XCTAssertEqual(expanded, ["org0", "org1"])   // two busiest only
    }

    func testHysteresisKeepsExpandedThenCollapses() {
        let policy = ExpansionPolicy(ttlTicks: 3)
        // Tick 1: busy -> expands.
        XCTAssertTrue(policy.update([D(key: "x", bytes: 1_000_000, endpointCount: 4)]).contains("x"))
        // Now idle, but still has endpoints: stays up for the TTL tail.
        XCTAssertTrue(policy.update([D(key: "x", bytes: 0, endpointCount: 4)]).contains("x"))   // ttl 2
        XCTAssertTrue(policy.update([D(key: "x", bytes: 0, endpointCount: 4)]).contains("x"))   // ttl 1
        XCTAssertFalse(policy.update([D(key: "x", bytes: 0, endpointCount: 4)]).contains("x"))  // ttl 0 -> collapsed
    }

    func testExpiredTimerWithLostEndpointsStaysCollapsed() {
        let policy = ExpansionPolicy(ttlTicks: 2)
        _ = policy.update([D(key: "x", bytes: 1_000_000, endpointCount: 4)])
        // Endpoints drop to 1 while timer still alive -> nothing to fan out.
        XCTAssertFalse(policy.update([D(key: "x", bytes: 0, endpointCount: 1)]).contains("x"))
    }
}