// Package u2 提供AES加密解密功能，支持ECB和CBC模式，以及PKCS5/PKCS7填充
package u2

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"errors"
)

// PaddingType 定义填充类型常量
type PaddingType int

const (
	// PaddingPKCS5 表示使用PKCS5填充，固定块大小为8字节
	PaddingPKCS5 PaddingType = iota

	// PaddingPKCS7 表示使用PKCS7填充，块大小可变（1-255字节）
	PaddingPKCS7
)

// pad 根据指定的填充类型对数据进行填充
//
// 参数:
//   - ciphertext: 需要填充的原始数据
//   - blockSize: 加密算法的块大小
//   - paddingType: 填充类型（PKCS5或PKCS7）
//
// 返回值:
//   - []byte: 填充后的数据
func pad(ciphertext []byte, blockSize int, paddingType PaddingType) []byte {
	var targetBlockSize int
	switch paddingType {
	case PaddingPKCS5:
		if blockSize != 8 {
			panic("PKCS5 padding requires 8-byte block size")
		}
		targetBlockSize = 8
	case PaddingPKCS7:
		targetBlockSize = blockSize
	default:
		targetBlockSize = blockSize
	}
	padding := targetBlockSize - len(ciphertext)%targetBlockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

// unpad 根据指定的填充类型去除数据中的填充内容
//
// 参数:
//   - origData: 需要去除填充的数据
//   - paddingType: 填充类型（PKCS5或PKCS7）
//
// 返回值:
//   - []byte: 去除填充后的原始数据
func unpad(origData []byte, paddingType PaddingType) []byte {
	length := len(origData)
	if length == 0 {
		return origData
	}
	unpadding := int(origData[length-1])

	// 验证unpadding的有效性
	switch paddingType {
	case PaddingPKCS5:
		if unpadding < 1 || unpadding > 8 {
			return origData
		}
		// 验证所有填充字节是否相同
		for i := length - unpadding; i < length; i++ {
			if int(origData[i]) != unpadding {
				return origData
			}
		}
	case PaddingPKCS7:
		if unpadding < 1 || unpadding > length {
			return origData
		}
		// 验证所有填充字节是否相同
		for i := length - unpadding; i < length; i++ {
			if int(origData[i]) != unpadding {
				return origData
			}
		}
	}
	return origData[:(length - unpadding)]
}

// validateKey 校验密钥长度，确保密钥长度为16、24或32字节
//
// 参数:
//   - key: 待校验的密钥
//
// 返回值:
//   - error: 如果密钥长度有效返回nil，否则返回错误
func validateKey(key []byte) error {
	switch len(key) {
	case 16, 24, 32:
		return nil
	default:
		return errors.New("AES密钥长度必须为128、192或256位（16、24、32字节）")
	}
}

// AesEncryptECBWithPadding 使用AES-ECB模式加密数据，并支持选择填充方式
//
// 参数:
//   - s: 需要加密的字符串
//   - key: 加密密钥
//   - paddingType: 填充类型（PKCS5或PKCS7）
//
// 返回值:
//   - string: 加密后的Base64编码字符串
//   - error: 加密过程中发生的错误
func AesEncryptECBWithPadding(s string, key string, paddingType PaddingType) (string, error) {
	if err := validateKey([]byte(key)); err != nil {
		return "", err
	}

	block, err := aes.NewCipher([]byte(key))
	if err != nil {
		return "", err
	}

	origData := []byte(s)
	blockSize := block.BlockSize()
	origData = pad(origData, blockSize, paddingType)
	encrypted := make([]byte, len(origData))

	// ECB模式加密
	for bs, be := 0, blockSize; bs < len(origData); bs, be = bs+blockSize, be+blockSize {
		block.Encrypt(encrypted[bs:be], origData[bs:be])
	}

	return base64.StdEncoding.EncodeToString(encrypted), nil
}

// AesDecryptECBWithPadding 使用AES-ECB模式解密数据，并支持选择填充方式
//
// 参数:
//   - s: 需要解密的Base64编码字符串
//   - key: 解密密钥
//   - paddingType: 填充类型（PKCS5或PKCS7）
//
// 返回值:
//   - string: 解密后的原始字符串
//   - error: 解密过程中发生的错误
func AesDecryptECBWithPadding(s string, key string, paddingType PaddingType) (string, error) {
	if err := validateKey([]byte(key)); err != nil {
		return "", err
	}

	src, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return "", err
	}

	block, err := aes.NewCipher([]byte(key))
	if err != nil {
		return "", err
	}

	blockSize := block.BlockSize()
	decrypted := make([]byte, len(src))

	// ECB模式解密
	for bs, be := 0, blockSize; bs < len(src); bs, be = bs+blockSize, be+blockSize {
		block.Decrypt(decrypted[bs:be], src[bs:be])
	}

	decrypted = unpad(decrypted, paddingType)
	return string(decrypted), nil
}

// AesEncryptCBCWithPadding 使用AES-CBC模式加密数据，并支持选择填充方式
//
// 参数:
//   - s: 需要加密的字符串
//   - key: 加密密钥
//   - vi: 初始化向量
//   - paddingType: 填充类型（PKCS5或PKCS7）
//
// 返回值:
//   - string: 加密后的Base64编码字符串
//   - error: 加密过程中发生的错误
func AesEncryptCBCWithPadding(s string, key string, vi string, paddingType PaddingType) (string, error) {
	if err := validateKey([]byte(key)); err != nil {
		return "", err
	}
	ekey, err := aes.NewCipher([]byte(key))
	if err != nil {
		return "", err
	}

	blockSize := ekey.BlockSize()
	origData := []byte(s)
	origData = pad(origData, blockSize, paddingType)
	blockMode := cipher.NewCBCEncrypter(ekey, []byte(vi))
	encrypted := make([]byte, len(origData))
	blockMode.CryptBlocks(encrypted, origData)
	return base64.StdEncoding.EncodeToString(encrypted), nil
}

// AesDecryptCBCWithPadding 使用AES-CBC模式解密数据，并支持选择填充方式
//
// 参数:
//   - s: 需要解密的Base64编码字符串
//   - key: 解密密钥
//   - vi: 初始化向量
//   - paddingType: 填充类型（PKCS5或PKCS7）
//
// 返回值:
//   - string: 解密后的原始字符串
//   - error: 解密过程中发生的错误
func AesDecryptCBCWithPadding(s string, key string, vi string, paddingType PaddingType) (string, error) {
	if err := validateKey([]byte(key)); err != nil {
		return "", err
	}
	src, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return "", err
	}
	block, err := aes.NewCipher([]byte(key))
	if err != nil {
		return "", err
	}
	blockMode := cipher.NewCBCDecrypter(block, []byte(vi))
	decrypted := make([]byte, len(src))
	blockMode.CryptBlocks(decrypted, src)
	decrypted = unpad(decrypted, paddingType)
	return string(decrypted), nil
}

// AesEncryptECB 使用AES-ECB模式加密数据，默认使用PKCS7填充
//
// 参数:
//   - s: 需要加密的字符串
//   - key: 加密密钥
//
// 返回值:
//   - string: 加密后的Base64编码字符串
//   - error: 加密过程中发生的错误
func AesEncryptECB(s string, key string) (string, error) {
	return AesEncryptECBWithPadding(s, key, PaddingPKCS7)
}

// AesDecryptECB 使用AES-ECB模式解密数据，默认使用PKCS7填充
//
// 参数:
//   - s: 需要解密的Base64编码字符串
//   - key: 解密密钥
//
// 返回值:
//   - string: 解密后的原始字符串
//   - error: 解密过程中发生的错误
func AesDecryptECB(s string, key string) (string, error) {
	return AesDecryptECBWithPadding(s, key, PaddingPKCS7)
}

// AesEncryptCBC 使用AES-CBC模式加密数据，默认使用PKCS7填充
//
// 参数:
//   - s: 需要加密的字符串
//   - key: 加密密钥
//   - vi: 初始化向量
//
// 返回值:
//   - string: 加密后的Base64编码字符串
//   - error: 加密过程中发生的错误
func AesEncryptCBC(s string, key string, vi string) (string, error) {
	return AesEncryptCBCWithPadding(s, key, vi, PaddingPKCS7)
}

// AesDecryptCBC 使用AES-CBC模式解密数据，默认使用PKCS7填充
//
// 参数:
//   - s: 需要解密的Base64编码字符串
//   - key: 解密密钥
//   - vi: 初始化向量
//
// 返回值:
//   - string: 解密后的原始字符串
//   - error: 解密过程中发生的错误
func AesDecryptCBC(s string, key string, vi string) (string, error) {
	return AesDecryptCBCWithPadding(s, key, vi, PaddingPKCS7)
}