import Foundation

// MARK: - KyberVariant

/**
 * Kyber KEM 安全级别枚举。
 *
 * CRYSTALS-Kyber（ML-KEM，NIST FIPS 203）定义了三种参数集，
 * 分别对应不同的安全强度和密钥/密文大小：
 *
 * | 变体         | KYBER_K | NIST 级别 | 经典安全性  |
 * |-------------|---------|----------|-----------|
 * | `.kyber512`  |    2    |  Level 1  |  ~AES-128 |
 * | `.kyber768`  |    3    |  Level 3  |  ~AES-192 |
 * | `.kyber1024` |    4    |  Level 5  |  ~AES-256 |
 *
 * 推荐选择：`.kyber768` 满足大多数应用场景的安全需求，
 * 同时密钥/密文大小适中。
 */
public enum KyberVariant: CaseIterable {
    /// Kyber512：NIST 安全级别 1，适用于资源受限或对大小敏感的场景
    case kyber512
    /// Kyber768：NIST 安全级别 3，推荐的通用安全级别
    case kyber768
    /// Kyber1024：NIST 安全级别 5，适用于高价值数据的长期保护
    case kyber1024

    /// 公钥字节长度
    public var publicKeyBytes: Int {
        switch self {
        case .kyber512:  return 800
        case .kyber768:  return 1184
        case .kyber1024: return 1568
        }
    }

    /// 私钥字节长度（包含公钥副本及 H(pk) 等辅助数据）
    public var secretKeyBytes: Int {
        switch self {
        case .kyber512:  return 1632
        case .kyber768:  return 2400
        case .kyber1024: return 3168
        }
    }

    /// 密文字节长度
    public var ciphertextBytes: Int {
        switch self {
        case .kyber512:  return 768
        case .kyber768:  return 1088
        case .kyber1024: return 1568
        }
    }

    /// 共享密钥字节长度（三种变体均为 32 字节）
    public var sharedSecretBytes: Int { 32 }

    /// 人类可读的展示名称
    public var displayName: String {
        switch self {
        case .kyber512:  return "Kyber512（NIST Level 1）"
        case .kyber768:  return "Kyber768（NIST Level 3）"
        case .kyber1024: return "Kyber1024（NIST Level 5）"
        }
    }
}

// MARK: - KyberKeyPair

/**
 * Kyber 密钥对，由 `KyberKEM.generateKeyPair` 返回。
 *
 * - `publicKey`：可安全分发给通信对方，对方使用它执行封装操作。
 * - `secretKey`：必须保密存储，用于从密文中恢复共享密钥。
 */
public struct KyberKeyPair {
    /// 公钥，可公开传输
    public let publicKey: Data
    /// 私钥，必须安全存储（建议使用 iOS Keychain）
    public let secretKey: Data

    public init(publicKey: Data, secretKey: Data) {
        self.publicKey = publicKey
        self.secretKey = secretKey
    }
}

// MARK: - KyberEncapsulationResult

/**
 * Kyber 封装结果，由 `KyberKEM.encapsulate` 返回。
 *
 * 封装操作（发送方/主动方）的两个输出：
 * - `ciphertext`：发送给持有私钥的接收方，接收方用私钥解封装以恢复共享密钥。
 * - `sharedSecret`：本地保留的 32 字节共享密钥，
 *                   可直接用作对称密码（AES-GCM / ChaCha20-Poly1305）的密钥材料。
 */
public struct KyberEncapsulationResult {
    /// 发送给接收方的密文
    public let ciphertext: Data
    /// 本地保留的共享密钥（32 字节）
    public let sharedSecret: Data

    public init(ciphertext: Data, sharedSecret: Data) {
        self.ciphertext = ciphertext
        self.sharedSecret = sharedSecret
    }
}

// MARK: - KyberError

/**
 * Kyber SDK 错误类型。
 *
 * 输入校验类错误在调用 C 层前抛出；
 * `*Failed` 类错误表示 C 函数返回了非零状态码（通常不应发生）。
 */
public enum KyberError: Error, LocalizedError {
    /// 密钥生成失败，附带 C 函数返回码
    case keyGenerationFailed(code: Int32)
    /// 封装操作失败，附带 C 函数返回码
    case encapsulationFailed(code: Int32)
    /// 解封装操作失败，附带 C 函数返回码
    case decapsulationFailed(code: Int32)
    /// 公钥长度与所选变体不符
    case invalidPublicKeySize(expected: Int, got: Int)
    /// 私钥长度与所选变体不符
    case invalidSecretKeySize(expected: Int, got: Int)
    /// 密文长度与所选变体不符
    case invalidCiphertextSize(expected: Int, got: Int)

    public var errorDescription: String? {
        switch self {
        case .keyGenerationFailed(let c):
            return "密钥生成失败（返回码：\(c)）"
        case .encapsulationFailed(let c):
            return "封装失败（返回码：\(c)）"
        case .decapsulationFailed(let c):
            return "解封装失败（返回码：\(c)）"
        case .invalidPublicKeySize(let e, let g):
            return "公钥长度不匹配：期望 \(e) 字节，实际 \(g) 字节"
        case .invalidSecretKeySize(let e, let g):
            return "私钥长度不匹配：期望 \(e) 字节，实际 \(g) 字节"
        case .invalidCiphertextSize(let e, let g):
            return "密文长度不匹配：期望 \(e) 字节，实际 \(g) 字节"
        }
    }
}