import XCTest
@testable import KyberSDK

/**
 * KyberSDK 单元测试套件
 *
 * 测试覆盖：
 *   - 三种变体（512 / 768 / 1024）密钥尺寸验证
 *   - 三种变体完整 KEM 全流程（密钥生成 → 封装 → 解封装 → 比对）
 *   - 错误密钥解封装（验证 IND-CCA2 安全属性）
 *   - 输入参数校验（非法公钥 / 私钥 / 密文长度）
 *   - Kyber768 性能基线
 */
final class KyberSDKTests: XCTestCase {

    // MARK: - Kyber512 测试

    /// 验证 Kyber512 密钥尺寸符合规范
    func testKyber512KeySizes() throws {
        let kp = try KyberKEM.generateKeyPair(variant: .kyber512)
        XCTAssertEqual(kp.publicKey.count, 800,  "Kyber512 公钥应为 800 字节")
        XCTAssertEqual(kp.secretKey.count, 1632, "Kyber512 私钥应为 1632 字节")
    }

    /// Kyber512 完整 KEM 流程：密钥生成 → 封装 → 解封装 → 共享密钥必须一致
    func testKyber512RoundTrip() throws {
        let kp  = try KyberKEM.generateKeyPair(variant: .kyber512)
        let enc = try KyberKEM.encapsulate(variant: .kyber512, publicKey: kp.publicKey)
        let dec = try KyberKEM.decapsulate(variant: .kyber512, ciphertext: enc.ciphertext, secretKey: kp.secretKey)

        XCTAssertEqual(enc.ciphertext.count,   768, "Kyber512 密文应为 768 字节")
        XCTAssertEqual(enc.sharedSecret.count, 32,  "共享密钥应为 32 字节")
        XCTAssertEqual(dec.count,              32,  "解封装共享密钥应为 32 字节")
        XCTAssertEqual(enc.sharedSecret, dec, "封装和解封装的共享密钥必须完全一致")
    }

    /// Kyber512 错误密钥测试：使用错误私钥解封装，共享密钥应与封装方不同（IND-CCA2）
    func testKyber512WrongKey() throws {
        let alice = try KyberKEM.generateKeyPair(variant: .kyber512)
        let bob   = try KyberKEM.generateKeyPair(variant: .kyber512)
        // 发送方使用 Alice 的公钥封装
        let enc   = try KyberKEM.encapsulate(variant: .kyber512, publicKey: alice.publicKey)
        // Bob 用自己的私钥解封装 —— 应得到不同的（伪随机）共享密钥
        let dec   = try KyberKEM.decapsulate(variant: .kyber512, ciphertext: enc.ciphertext, secretKey: bob.secretKey)
        XCTAssertNotEqual(enc.sharedSecret, dec, "使用错误私钥不应恢复正确的共享密钥（IND-CCA2）")
    }

    // MARK: - Kyber768 测试

    /// 验证 Kyber768 密钥尺寸符合规范
    func testKyber768KeySizes() throws {
        let kp = try KyberKEM.generateKeyPair(variant: .kyber768)
        XCTAssertEqual(kp.publicKey.count, 1184, "Kyber768 公钥应为 1184 字节")
        XCTAssertEqual(kp.secretKey.count, 2400, "Kyber768 私钥应为 2400 字节")
    }

    /// Kyber768 完整 KEM 流程
    func testKyber768RoundTrip() throws {
        let kp  = try KyberKEM.generateKeyPair(variant: .kyber768)
        let enc = try KyberKEM.encapsulate(variant: .kyber768, publicKey: kp.publicKey)
        let dec = try KyberKEM.decapsulate(variant: .kyber768, ciphertext: enc.ciphertext, secretKey: kp.secretKey)

        XCTAssertEqual(enc.ciphertext.count, 1088, "Kyber768 密文应为 1088 字节")
        XCTAssertEqual(enc.sharedSecret, dec, "封装和解封装的共享密钥必须完全一致")
    }

    // MARK: - Kyber1024 测试

    /// 验证 Kyber1024 密钥尺寸符合规范
    func testKyber1024KeySizes() throws {
        let kp = try KyberKEM.generateKeyPair(variant: .kyber1024)
        XCTAssertEqual(kp.publicKey.count, 1568, "Kyber1024 公钥应为 1568 字节")
        XCTAssertEqual(kp.secretKey.count, 3168, "Kyber1024 私钥应为 3168 字节")
    }

    /// Kyber1024 完整 KEM 流程
    func testKyber1024RoundTrip() throws {
        let kp  = try KyberKEM.generateKeyPair(variant: .kyber1024)
        let enc = try KyberKEM.encapsulate(variant: .kyber1024, publicKey: kp.publicKey)
        let dec = try KyberKEM.decapsulate(variant: .kyber1024, ciphertext: enc.ciphertext, secretKey: kp.secretKey)

        XCTAssertEqual(enc.ciphertext.count, 1568, "Kyber1024 密文应为 1568 字节")
        XCTAssertEqual(enc.sharedSecret, dec, "封装和解封装的共享密钥必须完全一致")
    }

    // MARK: - 输入参数校验测试

    /// 公钥长度不符时应抛出 KyberError.invalidPublicKeySize
    func testInvalidPublicKeySize() {
        XCTAssertThrowsError(
            try KyberKEM.encapsulate(variant: .kyber768, publicKey: Data(repeating: 0, count: 100)),
            "公钥长度不匹配时应抛出错误"
        )
    }

    /// 私钥长度不符时应抛出 KyberError.invalidSecretKeySize
    func testInvalidSecretKeySize() throws {
        let kp  = try KyberKEM.generateKeyPair(variant: .kyber768)
        let enc = try KyberKEM.encapsulate(variant: .kyber768, publicKey: kp.publicKey)
        XCTAssertThrowsError(
            try KyberKEM.decapsulate(
                variant: .kyber768,
                ciphertext: enc.ciphertext,
                secretKey: Data(repeating: 0, count: 10)
            ),
            "私钥长度不匹配时应抛出错误"
        )
    }

    /// 密文长度不符时应抛出 KyberError.invalidCiphertextSize
    func testInvalidCiphertextSize() throws {
        let kp = try KyberKEM.generateKeyPair(variant: .kyber768)
        XCTAssertThrowsError(
            try KyberKEM.decapsulate(
                variant: .kyber768,
                ciphertext: Data(repeating: 0, count: 10),
                secretKey: kp.secretKey
            ),
            "密文长度不匹配时应抛出错误"
        )
    }

    // MARK: - 性能测试

    /// Kyber768 密钥生成性能基线（10 次迭代平均值）
    func testKyber768Performance() throws {
        measure {
            _ = try? KyberKEM.generateKeyPair(variant: .kyber768)
        }
    }
}