import CryptoKit
import Foundation

/// Call media E2EE helpers.
///
/// - Shared key parity with desktop/android: `X25519 -> HSalsa20(zeros16, rawDh32)`.
/// - Frame crypto parity: `XChaCha20 xor(frame, nonce, key)` where nonce is
///   derived from frame additional-data timestamp.
enum CallMediaCrypto {

    static let keyLength = 32
    static let nonceLength = 24

    static func deriveSharedKey(
        localPrivateKey: Curve25519.KeyAgreement.PrivateKey,
        peerPublicHex: String
    ) -> Data? {
        let peerRaw = Data(hexString: peerPublicHex)
        guard peerRaw.count == keyLength else { return nil }
        guard let peerPublicKey = try? Curve25519.KeyAgreement.PublicKey(rawRepresentation: peerRaw) else {
            return nil
        }

        guard let rawDh = try? localPrivateKey.sharedSecretFromKeyAgreement(with: peerPublicKey) else {
            return nil
        }

        let rawDhBytes = rawDh.withUnsafeBytes { Data($0) }
        guard rawDhBytes.count == keyLength else { return nil }

        guard let hsalsa = NativeCryptoBridge.naclSharedKey(fromRawDH: rawDhBytes) else {
            return nil
        }
        guard hsalsa.count == keyLength else { return nil }
        return hsalsa
    }

    static func xorFrame(
        _ frame: Data,
        key: Data,
        additionalData: Data?
    ) -> Data? {
        guard key.count >= keyLength else { return nil }
        let nonce = nonceFromAdditionalData(additionalData)
        let key32 = Data(key.prefix(keyLength))
        return NativeCryptoBridge.xChaCha20Xor(frame, key: key32, nonce: nonce)
    }

    /// Desktop/audio parity:
    /// - preferred path: 8-byte big-endian timestamp payload.
    /// - fallback path: RTP header (timestamp bytes 4...7).
    private static func nonceFromAdditionalData(_ additionalData: Data?) -> Data {
        var nonce = Data(repeating: 0, count: nonceLength)
        guard let additionalData, !additionalData.isEmpty else { return nonce }

        if additionalData.count == 8 {
            nonce.replaceSubrange(0..<8, with: additionalData)
            return nonce
        }

        if additionalData.count >= 12 {
            let version = (additionalData[additionalData.startIndex] >> 6) & 0x03
            if version == 2 {
                let timestamp = additionalData[(additionalData.startIndex + 4)..<(additionalData.startIndex + 8)]
                nonce.replaceSubrange(4..<8, with: timestamp)
                return nonce
            }
        }

        if additionalData.count >= 8 {
            nonce.replaceSubrange(0..<8, with: additionalData.prefix(8))
        }

        return nonce
    }
}
