import Foundation

/// Reorders RTP packets by sequence number before depacketization so that a
/// reordering link (Wi-Fi) is not read as loss. Delivery stays in order; only
/// aged-out or overflowing buffers release out of order, which the downstream
/// gap detection still handles for genuine loss.
///
/// Mirrors the C++ `RtpJitterBuffer` (same defaults and semantics) and the
/// Kotlin receiver.
final class JitterBuffer {
    private struct Entry {
        let timeNs: Int64
        let packet: RtpPacket
    }

    private let maxDepth: Int
    private let maxDelayNs: Int64
    private var buffer: [Int: Entry] = [:]
    private var nextExpected: Int?

    init(maxDepth: Int = 16, maxDelayMs: Int = 60) {
        self.maxDepth = maxDepth
        self.maxDelayNs = Int64(maxDelayMs) * 1_000_000
    }

    /// Insert one packet and return the packets now ready for in-order delivery.
    func push(_ packet: RtpPacket) -> [RtpPacket] {
        var released: [RtpPacket] = []
        let sequence = packet.header.sequenceNumber
        let now = Self.nowNanos()

        let expected0: Int
        if let e = nextExpected {
            expected0 = e
        } else {
            expected0 = sequence
            nextExpected = sequence
        }

        // Serial-number comparison: a distance >= 32768 means the packet is
        // older than what we already delivered (duplicate or straggler).
        let distance = (sequence - expected0 + 65536) % 65536
        if distance < 32768 {
            buffer[sequence] = Entry(timeNs: now, packet: packet)

            // Release the consecutive run from the expected sequence.
            var expected = expected0
            while let entry = buffer[expected] {
                released.append(entry.packet)
                buffer.removeValue(forKey: expected)
                expected = (expected + 1) & 0xFFFF
            }
            nextExpected = expected

            // A missing packet stalls the run: age out the backlog (or bound
            // the buffer) and release what is there in order, so genuine loss
            // reaches the depacketizer's gap detection rather than blocking.
            if !buffer.isEmpty {
                let keys = buffer.keys.sorted()
                let head = keys.first!
                let headAgeNs = now - buffer[head]!.timeNs
                if headAgeNs > maxDelayNs || buffer.count > maxDepth {
                    for key in keys {
                        released.append(buffer[key]!.packet)
                    }
                    nextExpected = (keys.last! + 1) & 0xFFFF
                    buffer.removeAll(keepingCapacity: true)
                }
            }
        }
        return released
    }

    /// Discard everything still buffered.
    func clear() {
        buffer.removeAll(keepingCapacity: true)
        nextExpected = nil
    }

    private static func nowNanos() -> Int64 {
        return Int64(DispatchTime.now().uptimeNanoseconds)
    }
}