import Foundation

/// Writes RTP packets for one outgoing stream (RFC 3550).
struct RTPStream {
    let ssrc: UInt32
    let payloadType: UInt8
    private(set) var sequence: UInt16 = .random(in: 0...UInt16.max)
    private(set) var packetCount: UInt32 = 0
    private(set) var octetCount: UInt32 = 0

    init(ssrc: UInt32, payloadType: UInt8) {
        self.ssrc = ssrc
        self.payloadType = payloadType
    }

    mutating func packet(timestamp: UInt32, marker: Bool, payload: [UInt8]) -> [UInt8] {
        var w = ByteWriter(capacity: 12 + payload.count)
        w.u8(0x80)
        w.u8((marker ? 0x80 : 0) | (payloadType & 0x7F))
        w.u16(sequence)
        w.u32(timestamp)
        w.u32(ssrc)
        w.append(payload)
        sequence &+= 1
        packetCount &+= 1
        octetCount &+= UInt32(payload.count)
        return w.bytes
    }
}

/// H.264 packetization (RFC 6184, packetization-mode 1): single NAL units, STAP-A and FU-A.
enum H264Packetizer {
    /// Largest RTP payload; keeps IP packets well below 1500 bytes after SRTP and IPv6 headers.
    static let maxPayload = 1100

    /// Splits one access unit (NAL units without start codes or length prefixes) into RTP payloads.
    /// Small leading parameter sets (SPS/PPS) are aggregated into one STAP-A.
    static func payloads(for nalUnits: [[UInt8]]) -> [[UInt8]] {
        var out: [[UInt8]] = []
        var aggregate: [[UInt8]] = []

        func flushAggregate() {
            if aggregate.count == 1 {
                out.append(aggregate[0])
            } else if aggregate.count > 1 {
                var w = ByteWriter()
                w.u8((aggregate.map { $0[0] & 0x60 }.max() ?? 0) | 24) // STAP-A, highest NRI
                for nal in aggregate { w.u16(UInt16(nal.count)); w.append(nal) }
                out.append(w.bytes)
            }
            aggregate = []
        }

        for nal in nalUnits where !nal.isEmpty {
            let type = nal[0] & 0x1F
            if type == 7 || type == 8 {
                let size = aggregate.reduce(1) { $0 + 2 + $1.count } + 2 + nal.count
                if size > maxPayload { flushAggregate() }
                aggregate.append(nal)
                continue
            }
            flushAggregate()
            if nal.count <= maxPayload {
                out.append(nal)
                continue
            }
            let indicator = (nal[0] & 0xE0) | 28 // FU-A
            var offset = 1
            while offset < nal.count {
                let end = min(offset + maxPayload - 2, nal.count)
                var header = type
                if offset == 1 { header |= 0x80 }
                if end == nal.count { header |= 0x40 }
                out.append([indicator, header] + nal[offset..<end])
                offset = end
            }
        }
        flushAggregate()
        return out
    }

    /// Splits an AVCC buffer (4-byte big-endian length prefixes, as from VideoToolbox) into NAL units.
    static func nalUnits(avcc: [UInt8]) -> [[UInt8]] {
        var out: [[UInt8]] = []
        var i = 0
        while i + 4 <= avcc.count {
            let length = Int(avcc[i]) << 24 | Int(avcc[i + 1]) << 16 | Int(avcc[i + 2]) << 8 | Int(avcc[i + 3])
            i += 4
            guard length > 0, i + length <= avcc.count else { break }
            out.append(Array(avcc[i..<i + length]))
            i += length
        }
        return out
    }
}

/// RTCP feedback the browser sends that the sender must act on.
enum RTCPFeedback {
    /// True if a (compound) RTCP packet asks for a key frame (PLI or FIR).
    static func requestsKeyFrame(_ packet: [UInt8]) -> Bool {
        var i = 0
        while i + 4 <= packet.count {
            let fmt = packet[i] & 0x1F
            let type = packet[i + 1]
            let length = ((Int(packet[i + 2]) << 8 | Int(packet[i + 3])) + 1) * 4
            if (type == 206 && fmt == 1) || (type == 206 && fmt == 4) { return true }
            guard length > 0 else { return false }
            i += length
        }
        return false
    }

    /// RFC 7983: RTP/RTCP start with version 2 (first byte 128…191); RTCP payload types are 192…223.
    static func isRTCP(_ packet: [UInt8]) -> Bool {
        packet.count >= 8 && packet[0] >= 128 && packet[0] <= 191 && packet[1] >= 192 && packet[1] <= 223
    }

    /// Sender report so the browser can map RTP timestamps to wall clock (RFC 3550 6.4.1).
    static func senderReport(ssrc: UInt32, rtpTimestamp: UInt32, packets: UInt32, octets: UInt32, now: Date = Date()) -> [UInt8] {
        let ntp = now.timeIntervalSince1970 + 2_208_988_800
        let seconds = UInt32(ntp)
        let fraction = UInt32((ntp - Double(seconds)) * 4_294_967_296)
        var w = ByteWriter(capacity: 28)
        w.u8(0x80)
        w.u8(200)
        w.u16(6)
        w.u32(ssrc)
        w.u32(seconds)
        w.u32(fraction)
        w.u32(rtpTimestamp)
        w.u32(packets)
        w.u32(octets)
        return w.bytes
    }
}