import Foundation
import CKyber

// MARK: - KyberKEM

/**
 * Kyber KEM（密钥封装机制）Swift 封装层。
 *
 * ## 算法简介
 * CRYSTALS-Kyber（ML-KEM，NIST FIPS 203）是基于格密码（Module-LWE）的
 * 后量子密钥封装机制，由 NIST 于 2024 年正式标准化，
 * 可抵御量子计算机对传统公钥密码（RSA / ECDH）的攻击。
 *
 * ## 典型使用流程
 * ```swift
 * // 接收方：生成密钥对，公钥可公开分发
 * let keyPair = try KyberKEM.generateKeyPair(variant: .kyber768)
 *
 * // 发送方：用接收方公钥封装，得到密文和共享密钥
 * let result = try KyberKEM.encapsulate(variant: .kyber768, publicKey: keyPair.publicKey)
 * // 将 result.ciphertext 发送给接收方
 * // 将 result.sharedSecret 作为对称密钥使用（如 AES-GCM）
 *
 * // 接收方：用私钥解封装，恢复共享密钥
 * let sharedSecret = try KyberKEM.decapsulate(
 *     variant: .kyber768,
 *     ciphertext: result.ciphertext,
 *     secretKey: keyPair.secretKey
 * )
 * // result.sharedSecret == sharedSecret ✓
 * ```
 *
 * ## 安全说明
 * - 底层 C 实现的 `crypto_kem_keypair` 含有日期校验逻辑，
 *   本层全程调用 `keypair_derand`，由 `randombytes()` 生成随机数，
 *   彻底规避该限制。
 * - 解封装在密文被篡改时仍返回 0，但 `sharedSecret` 为伪随机值，
 *   这是 IND-CCA2 安全方案的标准行为。
 *
 * ## 线程安全性
 * 所有方法均为无状态纯函数，可在任意队列并发调用。
 */
public enum KyberKEM {

    // MARK: - 密钥生成

    /**
     * 生成 Kyber 密钥对。
     *
     * 使用 `arc4random_buf` 生成 64 字节随机 coins（iOS/macOS CSPRNG），
     * 调用 `keypair_derand` 确定性地生成密钥对。
     *
     * - Parameter variant: 安全级别，见 ``KyberVariant``
     * - Returns: 包含公钥和私钥的 ``KyberKeyPair``
     * - Throws: ``KyberError/keyGenerationFailed(code:)``（极罕见）
     */
    public static func generateKeyPair(variant: KyberVariant) throws -> KyberKeyPair {
        var pk = [UInt8](repeating: 0, count: variant.publicKeyBytes)
        var sk = [UInt8](repeating: 0, count: variant.secretKeyBytes)

        // 生成 64 字节随机 coins：前 32 字节用于 IND-CPA 密钥，后 32 字节用于拒绝采样 z
        var coins = [UInt8](repeating: 0, count: 64)
        randombytes(&coins, 64)

        let ret: Int32
        switch variant {
        case .kyber512:
            ret = pqcrystals_kyber512_ref_keypair_derand(&pk, &sk, coins)
        case .kyber768:
            ret = pqcrystals_kyber768_ref_keypair_derand(&pk, &sk, coins)
        case .kyber1024:
            ret = pqcrystals_kyber1024_ref_keypair_derand(&pk, &sk, coins)
        }

        guard ret == 0 else { throw KyberError.keyGenerationFailed(code: ret) }
        return KyberKeyPair(publicKey: Data(pk), secretKey: Data(sk))
    }

    // MARK: - 封装

    /**
     * 使用接收方公钥封装共享密钥（发送方调用）。
     *
     * 生成 32 字节随机 coins，调用 `enc_derand` 确定性封装。
     * 返回的 `ciphertext` 应传输给持有私钥的接收方；
     * `sharedSecret` 本地保留，用于对称加密。
     *
     * - Parameters:
     *   - variant:   安全级别，必须与生成公钥时使用的变体一致
     *   - publicKey: 接收方公钥
     * - Returns: ``KyberEncapsulationResult``，包含密文和共享密钥
     * - Throws: ``KyberError/invalidPublicKeySize(expected:got:)`` 或
     *           ``KyberError/encapsulationFailed(code:)``
     */
    public static func encapsulate(
        variant: KyberVariant,
        publicKey: Data
    ) throws -> KyberEncapsulationResult {
        guard publicKey.count == variant.publicKeyBytes else {
            throw KyberError.invalidPublicKeySize(
                expected: variant.publicKeyBytes, got: publicKey.count)
        }

        var ct = [UInt8](repeating: 0, count: variant.ciphertextBytes)
        var ss = [UInt8](repeating: 0, count: variant.sharedSecretBytes)

        // 生成 32 字节随机 coins 用于确定性封装
        var coins = [UInt8](repeating: 0, count: 32)
        randombytes(&coins, 32)

        let pk = Array(publicKey)
        let ret: Int32
        switch variant {
        case .kyber512:
            ret = pqcrystals_kyber512_ref_enc_derand(&ct, &ss, pk, coins)
        case .kyber768:
            ret = pqcrystals_kyber768_ref_enc_derand(&ct, &ss, pk, coins)
        case .kyber1024:
            ret = pqcrystals_kyber1024_ref_enc_derand(&ct, &ss, pk, coins)
        }

        guard ret == 0 else { throw KyberError.encapsulationFailed(code: ret) }
        return KyberEncapsulationResult(ciphertext: Data(ct), sharedSecret: Data(ss))
    }

    // MARK: - 解封装

    /**
     * 使用私钥从密文中恢复共享密钥（接收方调用）。
     *
     * IND-CCA2 安全性保证：若密文被篡改，返回的共享密钥为伪随机值，
     * 与封装方的共享密钥不同，从而无法建立有效的加密信道。
     *
     * - Parameters:
     *   - variant:    安全级别，必须与原始密钥对的变体一致
     *   - ciphertext: 封装方产生的密文
     *   - secretKey:  接收方私钥
     * - Returns: 32 字节共享密钥（若密文合法，则与封装方的共享密钥相同）
     * - Throws: ``KyberError/invalidCiphertextSize(expected:got:)``、
     *           ``KyberError/invalidSecretKeySize(expected:got:)`` 或
     *           ``KyberError/decapsulationFailed(code:)``
     */
    public static func decapsulate(
        variant: KyberVariant,
        ciphertext: Data,
        secretKey: Data
    ) throws -> Data {
        guard ciphertext.count == variant.ciphertextBytes else {
            throw KyberError.invalidCiphertextSize(
                expected: variant.ciphertextBytes, got: ciphertext.count)
        }
        guard secretKey.count == variant.secretKeyBytes else {
            throw KyberError.invalidSecretKeySize(
                expected: variant.secretKeyBytes, got: secretKey.count)
        }

        var ss = [UInt8](repeating: 0, count: variant.sharedSecretBytes)
        let ct = Array(ciphertext)
        let sk = Array(secretKey)

        let ret: Int32
        switch variant {
        case .kyber512:
            ret = pqcrystals_kyber512_ref_dec(&ss, ct, sk)
        case .kyber768:
            ret = pqcrystals_kyber768_ref_dec(&ss, ct, sk)
        case .kyber1024:
            ret = pqcrystals_kyber1024_ref_dec(&ss, ct, sk)
        }

        guard ret == 0 else { throw KyberError.decapsulationFailed(code: ret) }
        return Data(ss)
    }
}