import XCTest
@testable import Rosetta

/// Tests that iOS crypto primitives produce byte-identical output
/// to Android/Desktop for call E2EE.
final class CallCryptoParityTests: XCTestCase {

    // MARK: - HSalsa20 (nacl.box.before parity)

    func testHSalsa20ProducesConsistentOutput() throws {
        // HSalsa20 with zero nonce should produce deterministic output from any key
        let key = Data(repeating: 0x42, count: 32)
        let result1 = NativeCryptoBridge.naclSharedKey(fromRawDH: key)
        let result2 = NativeCryptoBridge.naclSharedKey(fromRawDH: key)

        XCTAssertNotNil(result1)
        XCTAssertNotNil(result2)
        XCTAssertEqual(result1!.count, 32)
        XCTAssertEqual(result1, result2, "HSalsa20 must be deterministic")
    }

    func testHSalsa20OutputDiffersForDifferentKeys() throws {
        let key1 = Data(repeating: 0x01, count: 32)
        let key2 = Data(repeating: 0x02, count: 32)

        let result1 = NativeCryptoBridge.naclSharedKey(fromRawDH: key1)
        let result2 = NativeCryptoBridge.naclSharedKey(fromRawDH: key2)

        XCTAssertNotNil(result1)
        XCTAssertNotNil(result2)
        XCTAssertNotEqual(result1, result2, "Different keys must produce different shared secrets")
    }

    func testHSalsa20RejectsWrongKeyLength() {
        let shortKey = Data(repeating: 0x00, count: 16)
        let result = NativeCryptoBridge.naclSharedKey(fromRawDH: shortKey)
        XCTAssertNil(result, "HSalsa20 must reject keys != 32 bytes")
    }

    // MARK: - XChaCha20 stream cipher

    func testXChaCha20EncryptDecryptRoundTrip() throws {
        let key = Data(repeating: 0xAB, count: 32)
        let nonce = Data(repeating: 0xCD, count: 24)
        let plaintext = Data("Hello Rosetta E2EE call frame!".utf8)

        let encrypted = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonce)
        XCTAssertNotNil(encrypted)
        XCTAssertEqual(encrypted!.count, plaintext.count, "Encrypted size must match plaintext size")
        XCTAssertNotEqual(encrypted, plaintext, "Encrypted data must differ from plaintext")

        // XOR cipher: encrypt again = decrypt
        let decrypted = NativeCryptoBridge.xChaCha20Xor(encrypted!, key: key, nonce: nonce)
        XCTAssertNotNil(decrypted)
        XCTAssertEqual(decrypted, plaintext, "XChaCha20 XOR decrypt must recover original plaintext")
    }

    func testXChaCha20DifferentNonceProducesDifferentCiphertext() throws {
        let key = Data(repeating: 0x11, count: 32)
        let nonce1 = Data(repeating: 0x01, count: 24)
        let nonce2 = Data(repeating: 0x02, count: 24)
        let plaintext = Data(repeating: 0xFF, count: 64)

        let enc1 = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonce1)
        let enc2 = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonce2)

        XCTAssertNotNil(enc1)
        XCTAssertNotNil(enc2)
        XCTAssertNotEqual(enc1, enc2, "Different nonces must produce different ciphertext")
    }

    func testXChaCha20DifferentKeyProducesDifferentCiphertext() throws {
        let key1 = Data(repeating: 0x11, count: 32)
        let key2 = Data(repeating: 0x22, count: 32)
        let nonce = Data(repeating: 0x00, count: 24)
        let plaintext = Data(repeating: 0xFF, count: 64)

        let enc1 = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key1, nonce: nonce)
        let enc2 = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key2, nonce: nonce)

        XCTAssertNotNil(enc1)
        XCTAssertNotNil(enc2)
        XCTAssertNotEqual(enc1, enc2, "Different keys must produce different ciphertext")
    }

    func testXChaCha20OutputSizeMatchesInput() throws {
        let key = Data(repeating: 0x33, count: 32)
        let nonce = Data(repeating: 0x44, count: 24)

        for size in [0, 1, 16, 63, 64, 65, 128, 960, 1024] {
            let plaintext = Data(repeating: 0xAA, count: size)
            let encrypted = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonce)
            XCTAssertNotNil(encrypted, "XChaCha20 must handle size \(size)")
            XCTAssertEqual(encrypted!.count, size, "Output size must equal input size for \(size) bytes")
        }
    }

    func testXChaCha20RejectsWrongKeyOrNonceLength() {
        let plaintext = Data(repeating: 0x00, count: 16)

        let badKey = NativeCryptoBridge.xChaCha20Xor(plaintext, key: Data(repeating: 0, count: 16), nonce: Data(repeating: 0, count: 24))
        XCTAssertNil(badKey, "Must reject key != 32 bytes")

        let badNonce = NativeCryptoBridge.xChaCha20Xor(plaintext, key: Data(repeating: 0, count: 32), nonce: Data(repeating: 0, count: 12))
        XCTAssertNil(badNonce, "Must reject nonce != 24 bytes")
    }

    // MARK: - Nonce construction parity (Desktop BigInt encoding)

    /// Verifies that a nonce built from an 8-byte big-endian timestamp
    /// matches the Desktop audioE2EE.ts fillNonceFromTimestamp output.
    func testNonceFromTimestampMatchesDesktopFormat() throws {
        let key = Data(repeating: 0x55, count: 32)
        let plaintext = Data(repeating: 0xEE, count: 48)

        // Desktop encodes timestamp 12345 as big-endian 8 bytes:
        // [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x30, 0x39]
        // Then nonce[8:24] = zeros
        let timestamp: UInt64 = 12345
        var nonce = Data(repeating: 0, count: 24)
        nonce[0] = UInt8((timestamp >> 56) & 0xFF)
        nonce[1] = UInt8((timestamp >> 48) & 0xFF)
        nonce[2] = UInt8((timestamp >> 40) & 0xFF)
        nonce[3] = UInt8((timestamp >> 32) & 0xFF)
        nonce[4] = UInt8((timestamp >> 24) & 0xFF)
        nonce[5] = UInt8((timestamp >> 16) & 0xFF)
        nonce[6] = UInt8((timestamp >> 8) & 0xFF)
        nonce[7] = UInt8(timestamp & 0xFF)

        // Encrypt with manually constructed nonce
        let enc1 = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonce)
        XCTAssertNotNil(enc1)

        // Same nonce built differently (raw bytes = big-endian 12345)
        var nonce2 = Data(repeating: 0, count: 24)
        var tsBE = timestamp.bigEndian
        nonce2.replaceSubrange(0..<8, with: Data(bytes: &tsBE, count: 8))

        let enc2 = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonce2)
        XCTAssertNotNil(enc2)

        XCTAssertEqual(enc1, enc2, "Both nonce construction methods must produce identical ciphertext")
    }

    /// Verifies that 32-bit RTP timestamp at nonce[4:8] matches
    /// the 64-bit encoding when value fits in 32 bits.
    func testRtpTimestamp32MatchesDesktop64ForSmallValues() throws {
        let key = Data(repeating: 0x66, count: 32)
        let plaintext = Data(repeating: 0xDD, count: 32)

        let rtpTs: UInt32 = 48000 // 1 second at 48kHz

        // Desktop-style: 64-bit big-endian in nonce[0:8]
        var nonceDesktop = Data(repeating: 0, count: 24)
        let ts64 = UInt64(rtpTs)
        nonceDesktop[0] = UInt8((ts64 >> 56) & 0xFF) // 0
        nonceDesktop[1] = UInt8((ts64 >> 48) & 0xFF) // 0
        nonceDesktop[2] = UInt8((ts64 >> 40) & 0xFF) // 0
        nonceDesktop[3] = UInt8((ts64 >> 32) & 0xFF) // 0
        nonceDesktop[4] = UInt8((ts64 >> 24) & 0xFF) // 0
        nonceDesktop[5] = UInt8((ts64 >> 16) & 0xFF) // 0
        nonceDesktop[6] = UInt8((ts64 >> 8) & 0xFF)  // 0xBB
        nonceDesktop[7] = UInt8(ts64 & 0xFF)          // 0x80

        // RTP-style: 32-bit big-endian at nonce[4:8]
        var nonceRtp = Data(repeating: 0, count: 24)
        nonceRtp[4] = UInt8((rtpTs >> 24) & 0xFF) // 0
        nonceRtp[5] = UInt8((rtpTs >> 16) & 0xFF) // 0
        nonceRtp[6] = UInt8((rtpTs >> 8) & 0xFF)  // 0xBB
        nonceRtp[7] = UInt8(rtpTs & 0xFF)          // 0x80

        let encDesktop = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonceDesktop)
        let encRtp = NativeCryptoBridge.xChaCha20Xor(plaintext, key: key, nonce: nonceRtp)

        XCTAssertNotNil(encDesktop)
        XCTAssertNotNil(encRtp)
        XCTAssertEqual(encDesktop, encRtp,
                       "32-bit RTP timestamp at nonce[4:8] must match 64-bit encoding for values < 2^32")
    }

    // MARK: - Full E2EE round-trip (simulate iOS encrypt → Desktop decrypt)

    func testFullE2EERoundTripWithDesktopNonce() throws {
        // Simulate: iOS encrypts a frame, Desktop decrypts with same key and nonce
        let sharedKey = Data(repeating: 0x77, count: 32)

        // Simulated Opus frame (TOC byte 0xFC = CELT, 20ms, 1 frame)
        var opusFrame = Data([0xFC]) // TOC: config=31, code=0
        opusFrame.append(Data(repeating: 0x42, count: 80)) // payload

        // Desktop-style timestamp nonce
        let timestamp: UInt64 = 96000 // 2 seconds at 48kHz
        var nonce = Data(repeating: 0, count: 24)
        for i in 0..<8 {
            nonce[i] = UInt8((timestamp >> UInt64((7 - i) * 8)) & 0xFF)
        }

        // Encrypt (iOS side)
        let encrypted = NativeCryptoBridge.xChaCha20Xor(opusFrame, key: sharedKey, nonce: nonce)
        XCTAssertNotNil(encrypted)
        XCTAssertEqual(encrypted!.count, opusFrame.count)
        XCTAssertNotEqual(encrypted, opusFrame)

        // Decrypt (simulating Desktop side with same key+nonce)
        let decrypted = NativeCryptoBridge.xChaCha20Xor(encrypted!, key: sharedKey, nonce: nonce)
        XCTAssertNotNil(decrypted)
        XCTAssertEqual(decrypted, opusFrame, "Desktop must recover exact Opus frame")
    }

    // MARK: - HSalsa20 + XChaCha20 full chain

    func testFullKeyDerivationAndEncryptionChain() throws {
        // Simulate two peers with fixed "DH results"
        let rawDH = Data((0..<32).map { UInt8($0) })

        let sharedKey = NativeCryptoBridge.naclSharedKey(fromRawDH: rawDH)
        XCTAssertNotNil(sharedKey)
        XCTAssertEqual(sharedKey!.count, 32)

        let plaintext = Data(repeating: 0xAA, count: 160) // typical Opus frame
        let nonce = Data(repeating: 0, count: 24)

        let encrypted = NativeCryptoBridge.xChaCha20Xor(plaintext, key: sharedKey!, nonce: nonce)
        XCTAssertNotNil(encrypted)

        let decrypted = NativeCryptoBridge.xChaCha20Xor(encrypted!, key: sharedKey!, nonce: nonce)
        XCTAssertEqual(decrypted, plaintext, "Full chain: HSalsa20 key derivation + XChaCha20 must round-trip")
    }
}