import Foundation
import os

/// Manually advanced `Clock<Duration>` double (T-iOS-4). Drives
/// `ReconnectMachine` / `PingScheduler` / `SessionEngine` tests with ZERO
/// real-time waiting: code under test suspends in `sleep(until:tolerance:)`
/// and only resumes when the test calls `advance(by:)` past its deadline.
///
/// Deterministic hand-off pattern (no yield-loops, no real sleeps):
///
///     let task = Task { try await scheduler.run() }   // sleeps on this clock
///     await clock.waitForSleepers(count: 1)           // sleeper is parked
///     clock.advance(by: Tunables.pingInterval)        // fire it
///
/// Sendable: all state lives behind one `OSAllocatedUnfairLock`; continuations
/// are always resumed OUTSIDE the lock. Cancelling a sleeping task throws
/// `CancellationError` out of `sleep`, matching `ContinuousClock` semantics.
public final class FakeClock: Clock, Sendable {
    /// Instant = duration offset from the clock's zero epoch.
    public struct Instant: InstantProtocol, Sendable, Hashable, Comparable {
        public typealias Duration = Swift.Duration

        public let offset: Swift.Duration

        public init(offset: Swift.Duration = .zero) {
            self.offset = offset
        }

        public func advanced(by duration: Swift.Duration) -> Instant {
            Instant(offset: offset + duration)
        }

        public func duration(to other: Instant) -> Swift.Duration {
            other.offset - offset
        }

        public static func < (lhs: Instant, rhs: Instant) -> Bool {
            lhs.offset < rhs.offset
        }
    }

    private struct Sleeper {
        let deadline: Instant
        let continuation: CheckedContinuation<Void, any Error>
    }

    private struct Waiter {
        let targetCount: Int
        let continuation: CheckedContinuation<Void, Never>
    }

    private struct State {
        var now: Instant
        var sleepers: [UUID: Sleeper] = [:]
        var waiters: [UUID: Waiter] = [:]
    }

    private enum SleepRegistration {
        case alreadyDue
        case cancelled
        case sleeping(satisfiedWaiters: [Waiter])
    }

    private let state: OSAllocatedUnfairLock<State>

    public init(now: Instant = Instant()) {
        state = OSAllocatedUnfairLock(initialState: State(now: now))
    }

    // MARK: - Clock conformance

    public var now: Instant {
        state.withLock { $0.now }
    }

    public var minimumResolution: Swift.Duration { .zero }

    public func sleep(until deadline: Instant, tolerance: Swift.Duration? = nil) async throws {
        let id = UUID()
        try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
                registerSleeper(id: id, deadline: deadline, continuation: continuation)
            }
        } onCancel: {
            let sleeper = state.withLock { $0.sleepers.removeValue(forKey: id) }
            sleeper?.continuation.resume(throwing: CancellationError())
        }
    }

    // MARK: - Manual control (test side)

    /// Move time forward and synchronously resume every sleeper whose deadline
    /// is now due (in deadline order) — resumed outside the lock.
    public func advance(by duration: Swift.Duration) {
        precondition(duration >= .zero, "FakeClock cannot move backwards")
        let due: [Sleeper] = state.withLock { state in
            state.now = state.now.advanced(by: duration)
            let dueNow = state.sleepers.filter { $0.value.deadline <= state.now }
            for key in dueNow.keys {
                state.sleepers.removeValue(forKey: key)
            }
            return dueNow.values.sorted { $0.deadline < $1.deadline }
        }
        for sleeper in due {
            sleeper.continuation.resume()
        }
    }

    /// Number of tasks currently parked in `sleep` — for test assertions.
    public var pendingSleeperCount: Int {
        state.withLock { $0.sleepers.count }
    }

    /// Suspend until at least `count` sleepers are parked. This is the
    /// deterministic "task reached its sleep" barrier — never poll or
    /// real-sleep to wait for the code under test. Resumes immediately if
    /// already satisfied; on task cancellation it resumes without error.
    public func waitForSleepers(count: Int = 1) async {
        let id = UUID()
        await withTaskCancellationHandler {
            await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
                registerWaiter(id: id, targetCount: count, continuation: continuation)
            }
        } onCancel: {
            let waiter = state.withLock { $0.waiters.removeValue(forKey: id) }
            waiter?.continuation.resume()
        }
    }

    // MARK: - Internals (registration under the lock, resumption outside it)

    private func registerSleeper(
        id: UUID,
        deadline: Instant,
        continuation: CheckedContinuation<Void, any Error>
    ) {
        let registration: SleepRegistration = state.withLock { state in
            if Task.isCancelled { return .cancelled }
            guard deadline > state.now else { return .alreadyDue }
            state.sleepers[id] = Sleeper(deadline: deadline, continuation: continuation)
            return .sleeping(satisfiedWaiters: Self.takeSatisfiedWaiters(&state))
        }
        switch registration {
        case .alreadyDue:
            continuation.resume()
        case .cancelled:
            continuation.resume(throwing: CancellationError())
        case .sleeping(let satisfiedWaiters):
            for waiter in satisfiedWaiters {
                waiter.continuation.resume()
            }
        }
    }

    private func registerWaiter(
        id: UUID,
        targetCount: Int,
        continuation: CheckedContinuation<Void, Never>
    ) {
        let isAlreadySatisfied: Bool = state.withLock { state in
            if Task.isCancelled || state.sleepers.count >= targetCount { return true }
            state.waiters[id] = Waiter(targetCount: targetCount, continuation: continuation)
            return false
        }
        if isAlreadySatisfied {
            continuation.resume()
        }
    }

    private static func takeSatisfiedWaiters(_ state: inout State) -> [Waiter] {
        let satisfiedKeys = state.waiters
            .filter { $0.value.targetCount <= state.sleepers.count }
            .map { $0.key }
        return satisfiedKeys.compactMap { state.waiters.removeValue(forKey: $0) }
    }
}
