import CommonCrypto
import CryptoKit
import Foundation

/// AES-128 in ECB mode, used as the block function for counter mode (CryptoKit has no AES-CTR).
final class AESBlock: @unchecked Sendable {
    private var cryptor: CCCryptorRef?

    init(key: [UInt8]) {
        precondition(key.count == 16)
        CCCryptorCreate(CCOperation(kCCEncrypt), CCAlgorithm(kCCAlgorithmAES), CCOptions(kCCOptionECBMode),
                        key, key.count, nil, &cryptor)
    }

    deinit { if let cryptor { CCCryptorRelease(cryptor) } }

    /// Counter-mode keystream: AES(iv), AES(iv+1), … (counter in the last 16 bits, RFC 3711 4.1.1).
    func keystream(iv: [UInt8], count: Int) -> [UInt8] {
        let blocks = (count + 15) / 16
        var input = [UInt8](repeating: 0, count: blocks * 16)
        for b in 0..<blocks {
            var block = iv
            let counter = (UInt32(block[14]) << 8 | UInt32(block[15])) &+ UInt32(b)
            block[14] = UInt8((counter >> 8) & 0xFF)
            block[15] = UInt8(counter & 0xFF)
            input.replaceSubrange(b * 16..<(b + 1) * 16, with: block)
        }
        var output = [UInt8](repeating: 0, count: input.count)
        var moved = 0
        CCCryptorUpdate(cryptor, input, input.count, &output, output.count, &moved)
        return Array(output.prefix(count))
    }
}

/// SRTP/SRTCP with the profile SRTP_AES128_CM_HMAC_SHA1_80 (RFC 3711, RFC 5764).
/// One context per direction; this side only sends RTP and reads RTCP feedback.
final class SRTPContext: @unchecked Sendable {
    struct SessionKeys {
        let cipher: AESBlock
        let salt: [UInt8]
        let auth: SymmetricKey
    }

    static let tagLength = 10

    let rtp: SessionKeys
    let rtcp: SessionKeys
    private var rolloverCounters: [UInt32: (roc: UInt32, lastSeq: UInt16)] = [:]
    private var srtcpIndex: UInt32 = 0

    init(masterKey: [UInt8], masterSalt: [UInt8]) {
        rtp = Self.derive(masterKey: masterKey, masterSalt: masterSalt, labels: (0, 1, 2))
        rtcp = Self.derive(masterKey: masterKey, masterSalt: masterSalt, labels: (3, 4, 5))
    }

    /// Key derivation (RFC 3711 4.3) with key_derivation_rate 0.
    static func deriveKey(masterKey: [UInt8], masterSalt: [UInt8], label: UInt8, length: Int) -> [UInt8] {
        var x = masterSalt + [0, 0]
        x[7] ^= label
        return AESBlock(key: masterKey).keystream(iv: x, count: length)
    }

    private static func derive(masterKey: [UInt8], masterSalt: [UInt8], labels: (UInt8, UInt8, UInt8)) -> SessionKeys {
        SessionKeys(
            cipher: AESBlock(key: deriveKey(masterKey: masterKey, masterSalt: masterSalt, label: labels.0, length: 16)),
            salt: deriveKey(masterKey: masterKey, masterSalt: masterSalt, label: labels.2, length: 14),
            auth: SymmetricKey(data: deriveKey(masterKey: masterKey, masterSalt: masterSalt, label: labels.1, length: 20)))
    }

    /// IV = salt·2^16 ⊕ SSRC·2^64 ⊕ index·2^16.
    private static func iv(salt: [UInt8], ssrc: UInt32, index: UInt64) -> [UInt8] {
        var iv = salt + [0, 0]
        for i in 0..<4 { iv[4 + i] ^= UInt8((ssrc >> UInt32(24 - 8 * i)) & 0xFF) }
        for i in 0..<6 { iv[8 + i] ^= UInt8((index >> UInt64(40 - 8 * i)) & 0xFF) }
        return iv
    }

    /// Encrypts and authenticates one RTP packet. Packets per SSRC must be passed in sending order.
    func protectRTP(_ packet: [UInt8]) -> [UInt8] {
        guard packet.count >= 12 else { return packet }
        let headerLength = Self.rtpHeaderLength(packet)
        let seq = UInt16(packet[2]) << 8 | UInt16(packet[3])
        let ssrc = UInt32(packet[8]) << 24 | UInt32(packet[9]) << 16 | UInt32(packet[10]) << 8 | UInt32(packet[11])

        var state = rolloverCounters[ssrc] ?? (0, seq)
        if seq < state.lastSeq && state.lastSeq - seq > 0x8000 { state.roc &+= 1 }
        state.lastSeq = seq
        rolloverCounters[ssrc] = state

        let index = UInt64(state.roc) << 16 | UInt64(seq)
        var out = packet
        let payloadCount = packet.count - headerLength
        if payloadCount > 0 {
            let stream = rtp.cipher.keystream(iv: Self.iv(salt: rtp.salt, ssrc: ssrc, index: index), count: payloadCount)
            for i in 0..<payloadCount { out[headerLength + i] ^= stream[i] }
        }
        let roc = [UInt8(state.roc >> 24), UInt8((state.roc >> 16) & 0xFF), UInt8((state.roc >> 8) & 0xFF), UInt8(state.roc & 0xFF)]
        let tag = HMAC<Insecure.SHA1>.authenticationCode(for: out + roc, using: rtp.auth)
        return out + Array(tag).prefix(Self.tagLength)
    }

    /// Encrypts and authenticates one (compound) RTCP packet.
    func protectRTCP(_ packet: [UInt8]) -> [UInt8] {
        guard packet.count >= 8 else { return packet }
        srtcpIndex = (srtcpIndex + 1) & 0x7FFF_FFFF
        let ssrc = UInt32(packet[4]) << 24 | UInt32(packet[5]) << 16 | UInt32(packet[6]) << 8 | UInt32(packet[7])
        var out = packet
        let stream = rtcp.cipher.keystream(iv: Self.iv(salt: rtcp.salt, ssrc: ssrc, index: UInt64(srtcpIndex)), count: packet.count - 8)
        for i in 0..<stream.count { out[8 + i] ^= stream[i] }
        let e = srtcpIndex | 0x8000_0000
        out += [UInt8(e >> 24), UInt8((e >> 16) & 0xFF), UInt8((e >> 8) & 0xFF), UInt8(e & 0xFF)]
        let tag = HMAC<Insecure.SHA1>.authenticationCode(for: out, using: rtcp.auth)
        return out + Array(tag).prefix(Self.tagLength)
    }

    /// Verifies and decrypts one SRTCP packet; nil if the tag does not match.
    func unprotectRTCP(_ packet: [UInt8]) -> [UInt8]? {
        guard packet.count >= 8 + 4 + Self.tagLength else { return nil }
        let authEnd = packet.count - Self.tagLength
        let tag = Array(packet[authEnd...])
        let expected = Array(HMAC<Insecure.SHA1>.authenticationCode(for: packet[0..<authEnd], using: rtcp.auth)).prefix(Self.tagLength)
        guard constantTimeEqual(tag, Array(expected)) else { return nil }
        let e = UInt32(packet[authEnd - 4]) << 24 | UInt32(packet[authEnd - 3]) << 16 | UInt32(packet[authEnd - 2]) << 8 | UInt32(packet[authEnd - 1])
        var out = Array(packet[0..<authEnd - 4])
        guard e & 0x8000_0000 != 0 else { return out }
        let ssrc = UInt32(packet[4]) << 24 | UInt32(packet[5]) << 16 | UInt32(packet[6]) << 8 | UInt32(packet[7])
        let stream = rtcp.cipher.keystream(iv: Self.iv(salt: rtcp.salt, ssrc: ssrc, index: UInt64(e & 0x7FFF_FFFF)), count: out.count - 8)
        for i in 0..<stream.count { out[8 + i] ^= stream[i] }
        return out
    }

    static func rtpHeaderLength(_ packet: [UInt8]) -> Int {
        var length = 12 + Int(packet[0] & 0x0F) * 4
        if packet[0] & 0x10 != 0, packet.count >= length + 4 {
            length += 4 + (Int(packet[length + 2]) << 8 | Int(packet[length + 3])) * 4
        }
        return min(length, packet.count)
    }
}