// Package chachax 提供 ChaCha20-Poly1305 与 XChaCha20-Poly1305 两种 AEAD
// （带关联数据的认证加密）封装。
//
// ChaCha20-Poly1305 是 RFC 8439 定义的 AEAD，也是 TLS 1.3 的强制套件，在没有 AES
// 硬件加速（如部分 ARM、旧 x86、受限云环境）时是首选；XChaCha20-Poly1305 使用 24
// 字节 nonce，允许在同一条密钥下安全地随机生成 nonce。
//
// 本包只做输入校验、nonce 生成与格式封装，密码学运算全部交给
// golang.org/x/crypto/chacha20poly1305，不自行实现任何密码学原语，也不做常量时间
// 比较（认证由底层 Open 保证）。
//
// 安全注意事项：
//   - nonce 在同一条密钥下绝不能重复。ChaCha20-Poly1305 的 12 字节 nonce 若随机
//     生成，同一密钥下最多安全使用约 2^32 次；XChaCha20-Poly1305 的 24 字节 nonce
//     则可以放心随机生成。
//   - 关联数据（AAD）可以为 nil，但加密与解密两次调用必须使用完全一致的 AAD，
//     否则认证失败。
//   - 本包的错误信息绝不包含密钥、明文、nonce 或密文内容，可以安全记录到日志。
package chachax

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

	"golang.org/x/crypto/chacha20poly1305"
)

const (
	// KeySize 是 ChaCha20-Poly1305 与 XChaCha20-Poly1305 的密钥长度（字节），固定为 32。
	// 任何长度不等于该值的密钥都会被拒绝。
	KeySize = chacha20poly1305.KeySize
	// NonceSize 是 ChaCha20-Poly1305 的 nonce 长度（字节），固定为 12。
	//
	// 安全注意：12 字节 nonce 太短，无法在同一条密钥下无限次安全地随机生成；
	// 若需要为每条消息随机取 nonce 而不必担心碰撞，请改用 XChaCha20-Poly1305。
	NonceSize = chacha20poly1305.NonceSize
	// NonceSizeX 是 XChaCha20-Poly1305 的 nonce 长度（字节），固定为 24。
	//
	// 安全注意：24 字节 nonce 允许在同一条密钥下安全地随机生成，不必担心碰撞。
	NonceSizeX = chacha20poly1305.NonceSizeX
	// Overhead 是 Poly1305 认证标签的长度（字节），固定为 16。
	// 密文长度恒等于明文长度加该值。
	Overhead = chacha20poly1305.Overhead
)

// 哨兵错误。调用方可用 errors.Is 区分失败类别，而无需解析错误文本。
var (
	// ErrInvalidKeySize 表示密钥长度不等于 KeySize（32 字节）。
	ErrInvalidKeySize = errors.New("chachax: invalid key size")
	// ErrInvalidNonceSize 表示 nonce 长度与该模式的 NonceSize / NonceSizeX 不符。
	ErrInvalidNonceSize = errors.New("chachax: invalid nonce size")
	// ErrCiphertextTooShort 表示密文长度不足以容纳认证标签（或 sealed 数据不足以
	// 容纳 nonce 与标签）。
	ErrCiphertextTooShort = errors.New("chachax: ciphertext is too short")
	// ErrAuthenticationFailed 表示认证失败：密文、nonce 或 AAD 被篡改，或密钥错误。
	//
	// 出现该错误时**不得使用任何返回的明文**（本包在所有失败分支都返回 nil 明文）。
	ErrAuthenticationFailed = errors.New("chachax: authentication failed")
)

// Encrypt 使用 ChaCha20-Poly1305（12 字节 nonce）加密 plaintext，并用 aad 做关联数据认证。
//
// 参数：
//   - key：长度必须恰好为 KeySize（32 字节），否则返回错误。
//   - plaintext：明文，可以为 nil 或空切片，此时密文只包含认证标签。
//   - aad：关联数据，可以为 nil；不保密但必须认证，解密时必须传入完全相同的值。
//
// 返回值：
//   - ciphertext：格式为「密文 ‖ 认证标签」，长度恒为 len(plaintext)+Overhead；调用方后续可以安全地自行修改。
//   - nonce：本次加密随机生成的 NonceSize（12）字节 nonce。
//   - err：密钥长度错误或随机数源不可用时返回错误，此时 ciphertext 与 nonce 均为 nil。
//
// 安全注意：返回的 nonce 必须与密文一起保存/传输，且**同一条密钥下绝不能重复使用**
// 同一个 nonce；一旦重复，机密性与完整性同时失效。调用方需自行负责 nonce 的存储。
// 返回的 key、plaintext、aad 切片不会被修改，也不会被保留引用。
func Encrypt(key, plaintext, aad []byte) (ciphertext, nonce []byte, err error) {
	aead, err := newAEAD(key)
	if err != nil {
		return nil, nil, err
	}

	nonce = make([]byte, NonceSize)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, nil, fmt.Errorf("chachax: 生成 nonce 失败: %w", err)
	}

	return aead.Seal(nil, nonce, plaintext, aad), nonce, nil
}

// Decrypt 使用 ChaCha20-Poly1305（12 字节 nonce）解密 ciphertext 并校验认证标签。
//
// 参数：
//   - key：长度必须恰好为 KeySize（32 字节），否则返回错误。
//   - ciphertext：Encrypt 返回的「密文 ‖ 认证标签」，长度必须 >= Overhead（16 字节）。
//   - nonce：加密时使用的 NonceSize（12）字节 nonce，长度必须精确匹配。
//   - aad：关联数据，必须与加密时传入的值完全一致（nil 与空切片等价）。
//
// 返回值：认证通过时返回解密出的明文；任何失败都返回错误且明文为 nil。
//
// 安全注意：本函数绝不返回未经认证的明文——密文被篡改、标签不匹配、nonce 或 aad
// 不一致时一律返回错误。调用方必须把错误当作「数据不可信」处理，不要复用或记录
// 半成品明文。
func Decrypt(key, ciphertext, nonce, aad []byte) ([]byte, error) {
	aead, err := newAEAD(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, fmt.Errorf("%w: 期望 %d 字节，实际 %d 字节", ErrInvalidNonceSize, aead.NonceSize(), len(nonce))
	}
	if len(ciphertext) < Overhead {
		return nil, fmt.Errorf("%w: 至少需要 %d 字节（认证标签），实际 %d 字节", ErrCiphertextTooShort, Overhead, len(ciphertext))
	}

	plaintext, err := aead.Open(nil, nonce, ciphertext, aad)
	if err != nil {
		return nil, fmt.Errorf("%w: 密文已被篡改或密钥/nonce/AAD 不匹配: %w", ErrAuthenticationFailed, err)
	}
	return plaintext, nil
}

// Seal 是 Encrypt 的便利形式：使用 ChaCha20-Poly1305（12 字节 nonce）加密，并把
// nonce 前置到返回值中，得到便于单字段存储的 sealed 格式。
//
// 返回的 sealed 格式固定为：nonce（NonceSize 字节） ‖ 密文 ‖ 认证标签，
// 总长度恒为 NonceSize + len(plaintext) + Overhead。
//
// 参数与安全注意同 Encrypt：key 长度必须为 KeySize，aad 可以为 nil 但解密时必须一致，
// 随机 nonce 由本函数内部生成并随返回值一起携带。
func Seal(key, plaintext, aad []byte) ([]byte, error) {
	ciphertext, nonce, err := Encrypt(key, plaintext, aad)
	if err != nil {
		return nil, err
	}
	return append(nonce, ciphertext...), nil
}

// Open 解开 Seal 生成的 sealed 数据（格式固定为 nonce ‖ 密文 ‖ 认证标签）。
//
// 参数：
//   - key：长度必须恰好为 KeySize（32 字节），否则返回错误。
//   - sealed：Seal 的返回值，长度必须 >= NonceSize+Overhead（28 字节）；
//     长度不足时返回错误，绝不越界切片。
//   - aad：关联数据，必须与 Seal 时传入的值完全一致（nil 与空切片等价）。
//
// 返回值：认证通过时返回明文；任何失败都返回错误且明文为 nil。
//
// 安全注意：与 Decrypt 相同，认证失败绝不返回未经认证的明文。
func Open(key, sealed, aad []byte) ([]byte, error) {
	if len(sealed) < NonceSize+Overhead {
		return nil, fmt.Errorf("%w: sealed 数据至少需要 %d 字节（nonce %d + 标签 %d），实际 %d 字节",
			ErrCiphertextTooShort, NonceSize+Overhead, NonceSize, Overhead, len(sealed))
	}
	return Decrypt(key, sealed[NonceSize:], sealed[:NonceSize], aad)
}

// EncryptX 使用 XChaCha20-Poly1305（24 字节 nonce）加密 plaintext，并用 aad 做关联数据认证。
//
// 参数与返回值含义同 Encrypt，区别仅在于 nonce 长度是 NonceSizeX（24）字节，
// 且返回的 ciphertext 仍为「密文 ‖ 认证标签」。
//
// 安全注意：24 字节 nonce 允许在同一条密钥下安全地随机生成而无需担心碰撞，因此
// 在无法可靠维护 nonce 计数器时，应优先使用本函数而不是 Encrypt。但即便 nonce 空间
// 很大，也仍不允许人为重复使用同一个 nonce。
func EncryptX(key, plaintext, aad []byte) (ciphertext, nonce []byte, err error) {
	aead, err := newAEADX(key)
	if err != nil {
		return nil, nil, err
	}

	nonce = make([]byte, NonceSizeX)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, nil, fmt.Errorf("chachax: 生成 XChaCha20 nonce 失败: %w", err)
	}

	return aead.Seal(nil, nonce, plaintext, aad), nonce, nil
}

// DecryptX 使用 XChaCha20-Poly1305（24 字节 nonce）解密 ciphertext 并校验认证标签。
//
// 参数：
//   - key：长度必须恰好为 KeySize（32 字节），否则返回错误。
//   - ciphertext：EncryptX 返回的「密文 ‖ 认证标签」，长度必须 >= Overhead（16 字节）。
//   - nonce：长度必须恰好为 NonceSizeX（24 字节），否则返回错误。
//   - aad：必须与加密时传入的值完全一致（nil 与空切片等价）。
//
// 返回值：认证通过时返回明文；任何失败都返回错误且明文为 nil。
func DecryptX(key, ciphertext, nonce, aad []byte) ([]byte, error) {
	aead, err := newAEADX(key)
	if err != nil {
		return nil, err
	}
	if len(nonce) != aead.NonceSize() {
		return nil, fmt.Errorf("%w: XChaCha20 期望 %d 字节，实际 %d 字节", ErrInvalidNonceSize, aead.NonceSize(), len(nonce))
	}
	if len(ciphertext) < Overhead {
		return nil, fmt.Errorf("%w: 至少需要 %d 字节（认证标签），实际 %d 字节", ErrCiphertextTooShort, Overhead, len(ciphertext))
	}

	plaintext, err := aead.Open(nil, nonce, ciphertext, aad)
	if err != nil {
		return nil, fmt.Errorf("%w: 密文已被篡改或密钥/nonce/AAD 不匹配: %w", ErrAuthenticationFailed, err)
	}
	return plaintext, nil
}

// SealX 是 EncryptX 的便利形式：使用 XChaCha20-Poly1305（24 字节 nonce）加密，并把
// nonce 前置到返回值中。
//
// 返回的 sealed 格式固定为：nonce（NonceSizeX 字节） ‖ 密文 ‖ 认证标签，
// 总长度恒为 NonceSizeX + len(plaintext) + Overhead。
//
// 由于 nonce 随机生成且空间足够大，本函数适合「一次加密、一个字段存全部」的场景。
func SealX(key, plaintext, aad []byte) ([]byte, error) {
	ciphertext, nonce, err := EncryptX(key, plaintext, aad)
	if err != nil {
		return nil, err
	}
	return append(nonce, ciphertext...), nil
}

// OpenX 解开 SealX 生成的 sealed 数据（格式固定为 nonce ‖ 密文 ‖ 认证标签）。
//
// 参数：
//   - key：长度必须恰好为 KeySize（32 字节），否则返回错误。
//   - sealed：SealX 的返回值，长度必须 >= NonceSizeX+Overhead（40 字节）；
//     长度不足时返回错误，绝不越界切片。
//   - aad：关联数据，必须与 SealX 时传入的值完全一致（nil 与空切片等价）。
//
// 返回值：认证通过时返回明文；任何失败都返回错误且明文为 nil。
func OpenX(key, sealed, aad []byte) ([]byte, error) {
	if len(sealed) < NonceSizeX+Overhead {
		return nil, fmt.Errorf("%w: sealed 数据至少需要 %d 字节（nonce %d + 标签 %d），实际 %d 字节",
			ErrCiphertextTooShort, NonceSizeX+Overhead, NonceSizeX, Overhead, len(sealed))
	}
	return DecryptX(key, sealed[NonceSizeX:], sealed[:NonceSizeX], aad)
}

// GenerateKey 使用 crypto/rand 生成一条长度恰好为 KeySize（32 字节）的随机密钥。
//
// 返回值：新生成的密钥字节，每次调用都不同；随机数源不可用时返回错误且密钥为 nil。
//
// 安全注意：返回的切片就是密钥本身，请妥善保管——不要写进日志、错误信息或
// stdout/stderr，用完可以调用方自行清零。
func GenerateKey() ([]byte, error) {
	key := make([]byte, KeySize)
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return nil, fmt.Errorf("chachax: 生成密钥失败: %w", err)
	}
	return key, nil
}

// GenerateKeyHex 使用 crypto/rand 生成随机密钥并返回其小写十六进制编码。
//
// 返回值：长度固定为 2*KeySize（64）的小写十六进制字符串；随机数源不可用时返回错误。
//
// 安全注意：该字符串是密钥的等价表示（十六进制只是编码，不是哈希或口令派生），
// 拿到它就能解密，因此必须按密钥级别保护；如需重新作为密钥使用，请先
// hex.DecodeString 得到原始字节再交给本包的函数。
func GenerateKeyHex() (string, error) {
	key, err := GenerateKey()
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(key), nil
}

// newAEAD 构造 12 字节 nonce 的 ChaCha20-Poly1305 实例，并统一包装密钥长度错误。
func newAEAD(key []byte) (cipher.AEAD, error) {
	if len(key) != KeySize {
		return nil, fmt.Errorf("%w: ChaCha20-Poly1305 要求 %d 字节，实际 %d 字节", ErrInvalidKeySize, KeySize, len(key))
	}
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, fmt.Errorf("chachax: 创建 ChaCha20-Poly1305 实例失败: %w", err)
	}
	return aead, nil
}

// newAEADX 构造 24 字节 nonce 的 XChaCha20-Poly1305 实例，并统一包装密钥长度错误。
func newAEADX(key []byte) (cipher.AEAD, error) {
	if len(key) != KeySize {
		return nil, fmt.Errorf("%w: XChaCha20-Poly1305 要求 %d 字节，实际 %d 字节", ErrInvalidKeySize, KeySize, len(key))
	}
	aead, err := chacha20poly1305.NewX(key)
	if err != nil {
		return nil, fmt.Errorf("chachax: 创建 XChaCha20-Poly1305 实例失败: %w", err)
	}
	return aead, nil
}