import Foundation

// MARK: - XChaCha20Engine

/// XChaCha20-Poly1305 AEAD cipher implementation.
/// Matches the Android `MessageCrypto` XChaCha20 implementation for cross-platform compatibility.
enum XChaCha20Engine {

    nonisolated static let poly1305TagSize = 16
    nonisolated private static let isNativeCryptoDisabled = ProcessInfo.processInfo.environment["ROSETTA_DISABLE_NATIVE_CRYPTO"] == "1"

    // MARK: - XChaCha20-Poly1305 Decrypt

    /// Decrypts ciphertext+tag using XChaCha20-Poly1305.
    nonisolated static func decrypt(ciphertextWithTag: Data, key: Data, nonce: Data) throws -> Data {
        guard ciphertextWithTag.count >= poly1305TagSize else {
            throw CryptoError.invalidData("Ciphertext too short for Poly1305 tag")
        }
        guard key.count == 32, nonce.count == 24 else {
            throw CryptoError.invalidData("Key must be 32 bytes, nonce must be 24 bytes")
        }

        if !isNativeCryptoDisabled {
            if let native = NativeCryptoBridge.xChaCha20Poly1305Decrypt(
                ciphertextWithTag,
                key: key,
                nonce: nonce
            ) {
                return native
            }
        }

        let ciphertext = ciphertextWithTag[0..<(ciphertextWithTag.count - poly1305TagSize)]
        let tag = ciphertextWithTag[(ciphertextWithTag.count - poly1305TagSize)...]

        // Step 1: HChaCha20 — derive subkey from key + first 16 bytes of nonce
        let subkey = hchacha20(key: key, nonce: Data(nonce[0..<16]))

        // Step 2: Build ChaCha20 nonce: [0,0,0,0] + nonce[16..<24]
        var chacha20Nonce = Data(repeating: 0, count: 12)
        chacha20Nonce[4..<12] = nonce[16..<24]

        // Step 3: Generate Poly1305 key from first 64 bytes of keystream (counter=0)
        let poly1305Key = chacha20Block(key: subkey, nonce: chacha20Nonce, counter: 0)[0..<32]

        // Step 4: Verify Poly1305 tag
        let computedTag = Poly1305Engine.mac(data: Data(ciphertext), key: Data(poly1305Key))

        guard constantTimeEqual(Data(tag), computedTag) else {
            throw CryptoError.decryptionFailed
        }

        // Step 5: Decrypt with ChaCha20 (counter starts at 1)
        return chacha20Encrypt(
            data: Data(ciphertext), key: subkey, nonce: chacha20Nonce, initialCounter: 1
        )
    }

    // MARK: - XChaCha20-Poly1305 Encrypt

    /// Encrypts plaintext using XChaCha20-Poly1305.
    nonisolated static func encrypt(plaintext: Data, key: Data, nonce: Data) throws -> Data {
        guard key.count == 32, nonce.count == 24 else {
            throw CryptoError.invalidData("Key must be 32 bytes, nonce must be 24 bytes")
        }

        if !isNativeCryptoDisabled {
            if let native = NativeCryptoBridge.xChaCha20Poly1305Encrypt(
                plaintext,
                key: key,
                nonce: nonce
            ) {
                return native
            }
        }

        // Step 1: HChaCha20 — derive subkey
        let subkey = hchacha20(key: key, nonce: Data(nonce[0..<16]))

        // Step 2: Build ChaCha20 nonce
        var chacha20Nonce = Data(repeating: 0, count: 12)
        chacha20Nonce[4..<12] = nonce[16..<24]

        // Step 3: Generate Poly1305 key
        let poly1305Key = chacha20Block(key: subkey, nonce: chacha20Nonce, counter: 0)[0..<32]

        // Step 4: Encrypt with ChaCha20 (counter starts at 1)
        let ciphertext = chacha20Encrypt(
            data: plaintext, key: subkey, nonce: chacha20Nonce, initialCounter: 1
        )

        // Step 5: Compute Poly1305 tag
        let tag = Poly1305Engine.mac(data: ciphertext, key: Data(poly1305Key))

        return ciphertext + tag
    }
}

// MARK: - ChaCha20 Core

extension XChaCha20Engine {

    /// ChaCha20 quarter round.
    nonisolated static func quarterRound(_ state: inout [UInt32], _ a: Int, _ b: Int, _ c: Int, _ d: Int) {
        state[a] = state[a] &+ state[b]; state[d] ^= state[a]; state[d] = (state[d] << 16) | (state[d] >> 16)
        state[c] = state[c] &+ state[d]; state[b] ^= state[c]; state[b] = (state[b] << 12) | (state[b] >> 20)
        state[a] = state[a] &+ state[b]; state[d] ^= state[a]; state[d] = (state[d] << 8)  | (state[d] >> 24)
        state[c] = state[c] &+ state[d]; state[b] ^= state[c]; state[b] = (state[b] << 7)  | (state[b] >> 25)
    }

    /// Generates a 64-byte ChaCha20 block.
    nonisolated static func chacha20Block(key: Data, nonce: Data, counter: UInt32) -> Data {
        var state = [UInt32](repeating: 0, count: 16)

        // Constants: "expand 32-byte k"
        state[0] = 0x61707865; state[1] = 0x3320646e
        state[2] = 0x79622d32; state[3] = 0x6b206574

        // Key (8 × UInt32, little-endian)
        for i in 0..<8 {
            state[4 + i] = key.withUnsafeBytes { $0.load(fromByteOffset: i * 4, as: UInt32.self).littleEndian }
        }

        state[12] = counter

        // Nonce (3 × UInt32, little-endian)
        for i in 0..<3 {
            state[13 + i] = nonce.withUnsafeBytes { $0.load(fromByteOffset: i * 4, as: UInt32.self).littleEndian }
        }

        var working = state

        // 20 rounds (10 double rounds)
        for _ in 0..<10 {
            quarterRound(&working, 0, 4, 8, 12);  quarterRound(&working, 1, 5, 9, 13)
            quarterRound(&working, 2, 6, 10, 14); quarterRound(&working, 3, 7, 11, 15)
            quarterRound(&working, 0, 5, 10, 15); quarterRound(&working, 1, 6, 11, 12)
            quarterRound(&working, 2, 7, 8, 13);  quarterRound(&working, 3, 4, 9, 14)
        }

        // Add initial state
        for i in 0..<16 { working[i] = working[i] &+ state[i] }

        // Serialize to bytes (little-endian)
        var result = Data(count: 64)
        for i in 0..<16 {
            let val = working[i].littleEndian
            result[i * 4]     = UInt8(truncatingIfNeeded: val)
            result[i * 4 + 1] = UInt8(truncatingIfNeeded: val >> 8)
            result[i * 4 + 2] = UInt8(truncatingIfNeeded: val >> 16)
            result[i * 4 + 3] = UInt8(truncatingIfNeeded: val >> 24)
        }
        return result
    }

    /// ChaCha20 stream cipher encryption/decryption.
    nonisolated static func chacha20Encrypt(data: Data, key: Data, nonce: Data, initialCounter: UInt32) -> Data {
        var result = Data(count: data.count)
        var counter = initialCounter

        for offset in stride(from: 0, to: data.count, by: 64) {
            let block = chacha20Block(key: key, nonce: nonce, counter: counter)
            let blockSize = min(64, data.count - offset)

            for i in 0..<blockSize {
                result[offset + i] = data[data.startIndex + offset + i] ^ block[i]
            }
            counter &+= 1
        }
        return result
    }

    /// HChaCha20: derives a 32-byte subkey from a 32-byte key and 16-byte nonce.
    nonisolated static func hchacha20(key: Data, nonce: Data) -> Data {
        var state = [UInt32](repeating: 0, count: 16)

        state[0] = 0x61707865; state[1] = 0x3320646e
        state[2] = 0x79622d32; state[3] = 0x6b206574

        for i in 0..<8 {
            state[4 + i] = key.withUnsafeBytes { $0.load(fromByteOffset: i * 4, as: UInt32.self).littleEndian }
        }
        for i in 0..<4 {
            state[12 + i] = nonce.withUnsafeBytes { $0.load(fromByteOffset: i * 4, as: UInt32.self).littleEndian }
        }

        for _ in 0..<10 {
            quarterRound(&state, 0, 4, 8, 12);  quarterRound(&state, 1, 5, 9, 13)
            quarterRound(&state, 2, 6, 10, 14); quarterRound(&state, 3, 7, 11, 15)
            quarterRound(&state, 0, 5, 10, 15); quarterRound(&state, 1, 6, 11, 12)
            quarterRound(&state, 2, 7, 8, 13);  quarterRound(&state, 3, 4, 9, 14)
        }

        // Output: first 4 words + last 4 words
        var result = Data(count: 32)
        for i in 0..<4 {
            let val = state[i].littleEndian
            result[i * 4]     = UInt8(truncatingIfNeeded: val)
            result[i * 4 + 1] = UInt8(truncatingIfNeeded: val >> 8)
            result[i * 4 + 2] = UInt8(truncatingIfNeeded: val >> 16)
            result[i * 4 + 3] = UInt8(truncatingIfNeeded: val >> 24)
        }
        for i in 0..<4 {
            let val = state[12 + i].littleEndian
            result[16 + i * 4]     = UInt8(truncatingIfNeeded: val)
            result[16 + i * 4 + 1] = UInt8(truncatingIfNeeded: val >> 8)
            result[16 + i * 4 + 2] = UInt8(truncatingIfNeeded: val >> 16)
            result[16 + i * 4 + 3] = UInt8(truncatingIfNeeded: val >> 24)
        }
        return result
    }

    /// Constant-time comparison of two Data objects.
    nonisolated static func constantTimeEqual(_ a: Data, _ b: Data) -> Bool {
        guard a.count == b.count else { return false }
        var result: UInt8 = 0
        for i in 0..<a.count {
            result |= a[a.startIndex + i] ^ b[b.startIndex + i]
        }
        return result == 0
    }
}
