// Package aesx 提供 AES 分组密码的加解密封装，支持 ECB/CBC/CFB/OFB/CTR 与 GCM。
//
// AES 是当前使用最广泛的对称分组密码：分组长度固定 16 字节，密钥长度可以是
// 16（AES-128）、24（AES-192）或 32（AES-256）字节。新代码建议使用 AES-256。
//
// # 模式选择建议
//
//   - AES-GCM：带认证的加密（AEAD），**首选**，同时保证机密性与完整性；
//   - AES-CTR / CFB / OFB：流式模式，**不提供完整性保护**，密文比特可被随意翻转；
//   - AES-CBC：需要填充，不提供完整性保护，且存在 padding oracle 风险；
//   - AES-ECB：**不具备语义安全性**，相同明文块会产生相同密文块，仅用于兼容旧系统。
//
// 除 GCM 之外的所有模式都必须另行配合消息认证码（例如 HMAC）使用，并采用
// 「先加密后认证」（Encrypt-then-MAC）的组合方式，否则无法发现密文篡改。
package aesx

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"

	"code.yun.ink/pkg/encryptx/cipherx"
)

const (
	// KeySize128 是 AES-128 的密钥长度（字节）。
	KeySize128 = 16
	// KeySize192 是 AES-192 的密钥长度（字节）。
	KeySize192 = 24
	// KeySize256 是 AES-256 的密钥长度（字节）。
	KeySize256 = 32
	// BlockSize 是 AES 的分组长度（字节）。
	BlockSize = aes.BlockSize
	// GCMNonceSize 是 GCM 推荐使用的 nonce 长度（字节）。
	//
	// 12 字节是与 GCM 设计匹配的「标准」长度：其他长度会先在内部做 GHASH 派生，
	// 既慢又更容易因实现差异导致互通问题。
	GCMNonceSize = 12
	// GCMTagSize 是 GCM 认证标签的长度（字节）。
	GCMTagSize = 16
)

// GCM 相关的错误定义。
var (
	// ErrInvalidNonceSize 表示 GCM 的 nonce 长度不等于 GCMNonceSize。
	ErrInvalidNonceSize = errors.New("aesx: invalid GCM nonce size")
	// ErrCiphertextTooShort 表示密文长度不足以容纳认证标签。
	ErrCiphertextTooShort = errors.New("aesx: ciphertext is shorter than the GCM tag")
	// ErrAuthenticationFailed 表示 GCM 认证失败：密文、nonce 或 AAD 被篡改，
	// 或密钥错误。解密结果不可信，调用方不得使用任何返回的明文。
	ErrAuthenticationFailed = errors.New("aesx: GCM authentication failed")
)

// Suite 绑定 AES 的分组构造方式与合法密钥长度，模式实现由 cipherx 提供。
//
// 导出它是为了让调用方能够以统一方式驱动任意模式（例如由配置决定模式时），
// 而无需在本包与 cipherx 之间写 switch。
var Suite = cipherx.NewSuite("AES", aes.NewCipher, BlockSize, KeySize128, KeySize192, KeySize256)

// ValidKeySize 报告 keySize 是否为合法的 AES 密钥长度。
func ValidKeySize(keySize int) bool {
	return keySize == KeySize128 || keySize == KeySize192 || keySize == KeySize256
}

// GenerateKey 用 crypto/rand 生成指定长度的 AES 密钥。
//
// keySize 必须是 16、24 或 32；调用方应妥善保存返回值，密钥一旦丢失或泄漏，
// 已加密的数据将无法保护。
func GenerateKey(keySize int) ([]byte, error) {
	if !ValidKeySize(keySize) {
		return nil, fmt.Errorf("aesx: 密钥长度 %d 非法, 允许 16/24/32: %w", keySize, cipherx.ErrInvalidKeySize)
	}
	key := make([]byte, keySize)
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return nil, fmt.Errorf("aesx: 生成密钥失败: %w", err)
	}
	return key, nil
}

// GenerateKeyHex 生成 AES 密钥并以十六进制字符串返回，便于写入配置文件。
//
// 注意：返回值是**编码后的文本**，不能直接当密钥使用。读回后必须先解码
// （encoding/hex 的 DecodeString，或本库的 hexx.Decode），再传给 Encrypt/Decrypt；
// 直接把字符串当密钥传入会因长度翻倍而被拒绝。
func GenerateKeyHex(keySize int) (string, error) {
	key, err := GenerateKey(keySize)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(key), nil
}

// ---- ECB ----
//
// 警告：ECB 不具备语义安全性，相同的明文分组会产生相同的密文分组，攻击者可以
// 借此识别明文中的重复结构与规律。仅为兼容既有系统而提供。

// EncryptECB 使用 AES-ECB 加密并做 PKCS#7 填充，返回大写十六进制字符串。
func EncryptECB(plaintext, key []byte) ([]byte, error) {
	return Suite.Encrypt(cipherx.ECB, plaintext, key)
}

// DecryptECB 解密 AES-ECB 密文。
func DecryptECB(ciphertext, key []byte) ([]byte, error) {
	return Suite.Decrypt(cipherx.ECB, ciphertext, key)
}

// EncryptECBHex 与 EncryptECB 相同，返回大写十六进制字符串。
func EncryptECBHex(plaintext, key []byte) (string, error) {
	return Suite.EncryptHex(cipherx.ECB, plaintext, key)
}

// DecryptECBHex 解密十六进制 AES-ECB 密文，大小写均可。
func DecryptECBHex(ciphertextHex string, key []byte) ([]byte, error) {
	return Suite.DecryptHex(cipherx.ECB, ciphertextHex, key)
}

// ---- CBC ----

// EncryptCBC 使用 AES-CBC 加密，随机 IV 前置在密文之前。
func EncryptCBC(plaintext, key []byte) ([]byte, error) {
	return Suite.Encrypt(cipherx.CBC, plaintext, key)
}

// DecryptCBC 解密 AES-CBC 密文（IV 前置格式）。
func DecryptCBC(ciphertext, key []byte) ([]byte, error) {
	return Suite.Decrypt(cipherx.CBC, ciphertext, key)
}

// EncryptCBCHex 与 EncryptCBC 相同，返回大写十六进制字符串。
func EncryptCBCHex(plaintext, key []byte) (string, error) {
	return Suite.EncryptHex(cipherx.CBC, plaintext, key)
}

// DecryptCBCHex 解密十六进制 AES-CBC 密文。
func DecryptCBCHex(ciphertextHex string, key []byte) ([]byte, error) {
	return Suite.DecryptHex(cipherx.CBC, ciphertextHex, key)
}

// EncryptCBCWithIV 使用调用方指定的 IV 加密，返回值不含 IV。
//
// 仅用于与规定了固定 IV 的既有系统互操作；自行管理 IV 时必须保证 IV 随机且
// 不可预测，并且在同一密钥下绝不重用。
func EncryptCBCWithIV(plaintext, key, iv []byte) ([]byte, error) {
	return Suite.EncryptWithIV(cipherx.CBC, plaintext, key, iv)
}

// DecryptCBCWithIV 使用调用方指定的 IV 解密。
func DecryptCBCWithIV(ciphertext, key, iv []byte) ([]byte, error) {
	return Suite.DecryptWithIV(cipherx.CBC, ciphertext, key, iv)
}

// ---- CFB ----

// EncryptCFB 使用 AES-CFB 加密，随机 IV 前置在密文之前。
func EncryptCFB(plaintext, key []byte) ([]byte, error) {
	return Suite.Encrypt(cipherx.CFB, plaintext, key)
}

// DecryptCFB 解密 AES-CFB 密文（IV 前置格式）。
func DecryptCFB(ciphertext, key []byte) ([]byte, error) {
	return Suite.Decrypt(cipherx.CFB, ciphertext, key)
}

// EncryptCFBHex 与 EncryptCFB 相同，返回大写十六进制字符串。
func EncryptCFBHex(plaintext, key []byte) (string, error) {
	return Suite.EncryptHex(cipherx.CFB, plaintext, key)
}

// DecryptCFBHex 解密十六进制 AES-CFB 密文。
func DecryptCFBHex(ciphertextHex string, key []byte) ([]byte, error) {
	return Suite.DecryptHex(cipherx.CFB, ciphertextHex, key)
}

// EncryptCFBWithIV 使用调用方指定的 IV 加密，返回值不含 IV。
func EncryptCFBWithIV(plaintext, key, iv []byte) ([]byte, error) {
	return Suite.EncryptWithIV(cipherx.CFB, plaintext, key, iv)
}

// DecryptCFBWithIV 使用调用方指定的 IV 解密。
func DecryptCFBWithIV(ciphertext, key, iv []byte) ([]byte, error) {
	return Suite.DecryptWithIV(cipherx.CFB, ciphertext, key, iv)
}

// ---- OFB ----

// EncryptOFB 使用 AES-OFB 加密，随机 IV 前置在密文之前。
func EncryptOFB(plaintext, key []byte) ([]byte, error) {
	return Suite.Encrypt(cipherx.OFB, plaintext, key)
}

// DecryptOFB 解密 AES-OFB 密文（IV 前置格式）。
func DecryptOFB(ciphertext, key []byte) ([]byte, error) {
	return Suite.Decrypt(cipherx.OFB, ciphertext, key)
}

// EncryptOFBHex 与 EncryptOFB 相同，返回大写十六进制字符串。
func EncryptOFBHex(plaintext, key []byte) (string, error) {
	return Suite.EncryptHex(cipherx.OFB, plaintext, key)
}

// DecryptOFBHex 解密十六进制 AES-OFB 密文。
func DecryptOFBHex(ciphertextHex string, key []byte) ([]byte, error) {
	return Suite.DecryptHex(cipherx.OFB, ciphertextHex, key)
}

// EncryptOFBWithIV 使用调用方指定的 IV 加密，返回值不含 IV。
func EncryptOFBWithIV(plaintext, key, iv []byte) ([]byte, error) {
	return Suite.EncryptWithIV(cipherx.OFB, plaintext, key, iv)
}

// DecryptOFBWithIV 使用调用方指定的 IV 解密。
func DecryptOFBWithIV(ciphertext, key, iv []byte) ([]byte, error) {
	return Suite.DecryptWithIV(cipherx.OFB, ciphertext, key, iv)
}

// ---- CTR ----

// EncryptCTR 使用 AES-CTR 加密，随机 IV 前置在密文之前。
func EncryptCTR(plaintext, key []byte) ([]byte, error) {
	return Suite.Encrypt(cipherx.CTR, plaintext, key)
}

// DecryptCTR 解密 AES-CTR 密文（IV 前置格式）。
//
// CTR 是对合的：加密与解密是同一个运算，因此本函数与 EncryptCTR 等价。
func DecryptCTR(ciphertext, key []byte) ([]byte, error) {
	return Suite.Decrypt(cipherx.CTR, ciphertext, key)
}

// EncryptCTRHex 与 EncryptCTR 相同，返回大写十六进制字符串。
func EncryptCTRHex(plaintext, key []byte) (string, error) {
	return Suite.EncryptHex(cipherx.CTR, plaintext, key)
}

// DecryptCTRHex 解密十六进制 AES-CTR 密文。
func DecryptCTRHex(ciphertextHex string, key []byte) ([]byte, error) {
	return Suite.DecryptHex(cipherx.CTR, ciphertextHex, key)
}

// EncryptCTRWithIV 使用调用方指定的 IV 加密，返回值不含 IV。
func EncryptCTRWithIV(plaintext, key, iv []byte) ([]byte, error) {
	return Suite.EncryptWithIV(cipherx.CTR, plaintext, key, iv)
}

// DecryptCTRWithIV 使用调用方指定的 IV 解密。
func DecryptCTRWithIV(ciphertext, key, iv []byte) ([]byte, error) {
	return Suite.DecryptWithIV(cipherx.CTR, ciphertext, key, iv)
}

// ---- GCM ----
//
// GCM 是本包唯一提供完整性保护的 AEAD 模式，应作为默认选择。

func newGCM(key []byte) (cipher.AEAD, error) {
	block, err := Suite.NewBlock(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("aesx: 创建 GCM 失败: %w", err)
	}
	return aead, nil
}

// NewGCMNonce 生成一个随机的 12 字节 GCM nonce。
//
// nonce 本身无需保密，但**在同一密钥下绝不能重复**：GCM 一旦重用 nonce，
// 攻击者即可恢复认证密钥并伪造密文。使用同一个密钥加密超过 2^32 条消息时
// 应更换密钥。
func NewGCMNonce() ([]byte, error) {
	nonce := make([]byte, GCMNonceSize)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("aesx: 生成 GCM nonce 失败: %w", err)
	}
	return nonce, nil
}

// EncryptGCM 使用 AES-GCM 加密，随机生成 nonce 并单独返回。
//
// aad（附加认证数据）会参与完整性校验但不被加密，可用于绑定协议头、版本号、
// 消息 ID 等上下文；不需要时传 nil。密文返回值已包含认证标签。
//
// 调用方必须把 nonce 与密文一起保存或传输。
func EncryptGCM(plaintext, key, aad []byte) (ciphertext, nonce []byte, err error) {
	aead, err := newGCM(key)
	if err != nil {
		return nil, nil, err
	}
	nonce, err = NewGCMNonce()
	if err != nil {
		return nil, nil, err
	}
	return aead.Seal(nil, nonce, plaintext, aad), nonce, nil
}

// EncryptGCMWithNonce 使用调用方提供的 nonce 加密，返回值不含 nonce。
//
// 仅用于与既有系统互操作。传入的 nonce 必须与密文一同保存，且在同一密钥下
// 绝不重复使用。
func EncryptGCMWithNonce(plaintext, key, nonce, aad []byte) ([]byte, error) {
	aead, err := newGCM(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, fmt.Errorf("aesx: GCM nonce 长度 %d, 期望 %d: %w",
			len(nonce), aead.NonceSize(), ErrInvalidNonceSize)
	}
	return aead.Seal(nil, nonce, plaintext, aad), nil
}

// DecryptGCM 使用 AES-GCM 解密并验证完整性。
//
// 认证失败时返回包装了 ErrAuthenticationFailed 的错误，此时**返回的明文为
// nil**，调用方绝不可使用任何部分解密结果。
func DecryptGCM(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	aead, err := newGCM(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, fmt.Errorf("aesx: GCM nonce 长度 %d, 期望 %d: %w",
			len(nonce), aead.NonceSize(), ErrInvalidNonceSize)
	}
	// 长度校验必须在 Open 之前：认证标签至少要占 Overhead 字节。
	if len(ciphertext) < aead.Overhead() {
		return nil, fmt.Errorf("aesx: GCM 密文长度 %d 小于认证标签长度 %d: %w",
			len(ciphertext), aead.Overhead(), ErrCiphertextTooShort)
	}
	plaintext, err := aead.Open(nil, nonce, ciphertext, aad)
	if err != nil {
		return nil, fmt.Errorf("%w: %w", ErrAuthenticationFailed, err)
	}
	return plaintext, nil
}

// SealGCM 加密并把 nonce 前置在密文中，返回 nonce||密文。
//
// 单字段存储场景（数据库一个列、一个 token）用这个函数最省事：nonce 随密文
// 一起存放，不会漏掉。
func SealGCM(plaintext, key, aad []byte) ([]byte, error) {
	ciphertext, nonce, err := EncryptGCM(plaintext, key, aad)
	if err != nil {
		return nil, err
	}
	sealed := make([]byte, 0, len(nonce)+len(ciphertext))
	sealed = append(sealed, nonce...)
	sealed = append(sealed, ciphertext...)
	return sealed, nil
}

// OpenGCM 解密 SealGCM 产生的数据（nonce||密文 格式）。
func OpenGCM(sealed, key, aad []byte) ([]byte, error) {
	if len(sealed) < GCMNonceSize+GCMTagSize {
		return nil, fmt.Errorf("aesx: 数据长度 %d 不足以容纳 nonce 与认证标签: %w",
			len(sealed), ErrCiphertextTooShort)
	}
	nonce := sealed[:GCMNonceSize]
	ciphertext := sealed[GCMNonceSize:]
	return DecryptGCM(ciphertext, key, nonce, aad)
}

// EncryptGCMHex 与 EncryptGCM 相同，但密文与 nonce 均以十六进制字符串返回。
func EncryptGCMHex(plaintext, key, aad []byte) (ciphertextHex, nonceHex string, err error) {
	ciphertext, nonce, err := EncryptGCM(plaintext, key, aad)
	if err != nil {
		return "", "", err
	}
	return hex.EncodeToString(ciphertext), hex.EncodeToString(nonce), nil
}

// DecryptGCMHex 解密十六进制形式的 GCM 密文与 nonce。
func DecryptGCMHex(ciphertextHex, nonceHex string, key, aad []byte) ([]byte, error) {
	ciphertext, err := hex.DecodeString(ciphertextHex)
	if err != nil {
		return nil, fmt.Errorf("aesx: 密文十六进制解码失败: %w", err)
	}
	nonce, err := hex.DecodeString(nonceHex)
	if err != nil {
		return nil, fmt.Errorf("aesx: nonce 十六进制解码失败: %w", err)
	}
	return DecryptGCM(ciphertext, key, nonce, aad)
}