import Foundation
import CryptoKit

/// Core cryptographic utilities: AES-GCM, HKDF, KDF helpers
enum CryptoUtils {

    // MARK: - AES-256-GCM

    /// Encrypt with AES-256-GCM. Returns (key, nonce, ciphertext, tag) — all as Data.
    /// If key is nil, generates a random 256-bit key.
    /// Matches Python: aes_encrypt(plaintext, key=None)
    static func aesEncrypt(_ plaintext: Data, key: Data? = nil) throws -> (key: Data, nonce: Data, ciphertext: Data, tag: Data) {
        let keyData = key ?? Data.randomBytes(32)
        let symmetricKey = SymmetricKey(data: keyData)
        let nonceData = Data.randomBytes(12)
        let gcmNonce = try AES.GCM.Nonce(data: nonceData)

        let sealedBox = try AES.GCM.seal(plaintext, using: symmetricKey, nonce: gcmNonce)

        return (
            key: keyData,
            nonce: nonceData,
            ciphertext: Data(sealedBox.ciphertext),
            tag: Data(sealedBox.tag)
        )
    }

    /// Decrypt with AES-256-GCM.
    /// Matches Python: aes_decrypt(key, nonce, ciphertext, tag)
    static func aesDecrypt(key: Data, nonce: Data, ciphertext: Data, tag: Data) throws -> Data {
        let symmetricKey = SymmetricKey(data: key)
        let gcmNonce = try AES.GCM.Nonce(data: nonce)

        let sealedBox = try AES.GCM.SealedBox(
            nonce: gcmNonce,
            ciphertext: ciphertext,
            tag: tag
        )

        do {
            return try AES.GCM.open(sealedBox, using: symmetricKey)
        } catch {
            throw CryptoError.decryptionFailed("AES-GCM decryption failed")
        }
    }

    /// Encrypt with AES-256-GCM using AAD. Returns ciphertext with tag appended.
    /// Used by Double Ratchet and Sender Keys.
    static func aesGcmEncrypt(_ plaintext: Data, key: Data, nonce: Data, aad: Data) throws -> Data {
        let symmetricKey = SymmetricKey(data: key)
        let gcmNonce = try AES.GCM.Nonce(data: nonce)

        let sealedBox = try AES.GCM.seal(
            plaintext,
            using: symmetricKey,
            nonce: gcmNonce,
            authenticating: aad
        )

        // Return ciphertext + tag concatenated (matches Python AESGCM.encrypt)
        return Data(sealedBox.ciphertext) + Data(sealedBox.tag)
    }

    /// Decrypt AES-256-GCM with AAD. Input ciphertext has tag appended (last 16 bytes).
    static func aesGcmDecrypt(_ ctWithTag: Data, key: Data, nonce: Data, aad: Data) throws -> Data {
        guard ctWithTag.count >= 16 else {
            throw CryptoError.decryptionFailed("Ciphertext too short")
        }

        let ct = ctWithTag.prefix(ctWithTag.count - 16)
        let tag = ctWithTag.suffix(16)

        let symmetricKey = SymmetricKey(data: key)
        let gcmNonce = try AES.GCM.Nonce(data: nonce)

        let sealedBox = try AES.GCM.SealedBox(
            nonce: gcmNonce,
            ciphertext: ct,
            tag: tag
        )

        do {
            return try AES.GCM.open(sealedBox, using: symmetricKey, authenticating: aad)
        } catch {
            throw CryptoError.decryptionFailed("AES-GCM decryption with AAD failed")
        }
    }

    // MARK: - HKDF

    /// HKDF-SHA256 key derivation.
    /// Matches Python: hkdf_derive(input_key, salt, info, length=32)
    static func hkdfDerive(inputKey: Data, salt: Data, info: Data, length: Int = 32) -> Data {
        let symmetricKey = SymmetricKey(data: inputKey)
        let derived = HKDF<SHA256>.deriveKey(
            inputKeyMaterial: symmetricKey,
            salt: salt,
            info: info,
            outputByteCount: length
        )
        return derived.withUnsafeBytes { Data($0) }
    }

    // MARK: - KDF for Double Ratchet

    /// Root key KDF. Returns (newRootKey, chainKey).
    /// HKDF with rootKey as salt and DH output as input. Derives 64 bytes, split in half.
    /// Matches Python: kdf_rk(root_key, dh_output)
    static func kdfRK(rootKey: Data, dhOutput: Data) -> (newRootKey: Data, chainKey: Data) {
        let derived = hkdfDerive(
            inputKey: dhOutput,
            salt: rootKey,
            info: Data(Constants.rootKeyInfo.utf8),
            length: 64
        )
        return (derived.prefix(32), Data(derived.suffix(32)))
    }

    /// Chain key KDF. Returns (newChainKey, messageKey).
    /// HMAC-SHA256: messageKey = HMAC(chainKey, 0x01), newChainKey = HMAC(chainKey, 0x02)
    /// Matches Python: kdf_ck(chain_key)
    static func kdfCK(chainKey: Data) -> (newChainKey: Data, messageKey: Data) {
        let symmetricKey = SymmetricKey(data: chainKey)
        let messageKey = Data(HMAC<SHA256>.authenticationCode(for: Data([0x01]), using: symmetricKey))
        let newChainKey = Data(HMAC<SHA256>.authenticationCode(for: Data([0x02]), using: symmetricKey))
        return (newChainKey, messageKey)
    }

    // MARK: - Self-Encryption Key

    /// Derive static AES-256 key from identity key for self-encrypted message copies.
    /// Matches Python: derive_self_encryption_key(identity_private)
    static func deriveSelfEncryptionKey(identityPrivateRaw: Data) -> Data {
        hkdfDerive(
            inputKey: identityPrivateRaw,
            salt: Data(Constants.selfEncryptionSalt.utf8),
            info: Data(Constants.selfEncryptionInfo.utf8),
            length: 32
        )
    }

    // MARK: - Local Storage Key

    /// Derive AES-256 key for encrypting local session/sender key files.
    /// Matches Python: derive_local_storage_key(identity_private)
    static func deriveLocalStorageKey(identityPrivateRaw: Data) -> Data {
        hkdfDerive(
            inputKey: identityPrivateRaw,
            salt: Data(Constants.localStorageSalt.utf8),
            info: Data(Constants.localStorageInfo.utf8),
            length: 32
        )
    }

    // MARK: - Local File Encryption

    /// Encrypt data for local storage. Format: nonce(12) + tag(16) + ciphertext
    /// Matches Python: _encrypt_local(data, key)
    static func encryptLocal(_ data: Data, key: Data) throws -> Data {
        let symmetricKey = SymmetricKey(data: key)
        let sealedBox = try AES.GCM.seal(data, using: symmetricKey)

        var result = Data()
        result.append(Data(sealedBox.nonce))       // 12 bytes
        result.append(Data(sealedBox.tag))          // 16 bytes
        result.append(Data(sealedBox.ciphertext))   // N bytes
        return result
    }

    /// Decrypt locally stored data. Format: nonce(12) + tag(16) + ciphertext
    /// Matches Python: _decrypt_local(raw, key)
    static func decryptLocal(_ raw: Data, key: Data) throws -> Data {
        guard raw.count >= 28 else {  // 12 + 16 minimum
            throw CryptoError.decryptionFailed("Local encrypted data too short")
        }

        let nonce = raw[0..<12]
        let tag = raw[12..<28]
        let ct = raw[28...]

        let symmetricKey = SymmetricKey(data: key)
        let gcmNonce = try AES.GCM.Nonce(data: nonce)

        let sealedBox = try AES.GCM.SealedBox(
            nonce: gcmNonce,
            ciphertext: ct,
            tag: tag
        )

        do {
            return try AES.GCM.open(sealedBox, using: symmetricKey)
        } catch {
            throw CryptoError.decryptionFailed("Local storage decryption failed")
        }
    }
}