package xlog

import (
	"crypto/aes"
	"crypto/cipher"
	"errors"
	"fmt"
)

// 外部禁止使用此结构体创建实例，否则会出错
type XlogEncrypt struct {
	AesIv []byte // 长度必须为aes.BlockSize
}

var aes_key = []byte("12369854712564897564325467553685") // 32 字节 = 256 位
var aes_iv = []byte("5123688955621365")

// 必须用函数创建实例，否则会出错
func NewXlogEncrypt() XlogEncrypt {
	xe := XlogEncrypt{
		AesIv: aes_iv,
	}
	return xe
}

func (x *XlogEncrypt) SetNewIv(newIv []byte) error {
	if len(newIv) != aes.BlockSize {
		return fmt.Errorf("IV length must be 16 bytes, got %d", len(newIv))
	}

	x.AesIv = newIv
	return nil
}

// PKCS7 填充
func PKCS7Pad(data []byte, blockSize int) []byte {
	padLen := blockSize - len(data)%blockSize
	pad := make([]byte, padLen)
	for i := range pad {
		pad[i] = byte(padLen)
	}
	return append(data, pad...)
}

// PKCS7 去填充
func PKCS7Unpad(data []byte, blockSize int) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, errors.New("data is empty")
	}
	if length%blockSize != 0 {
		return nil, errors.New("data is not a multiple of block size")
	}
	padLen := int(data[length-1])
	if padLen == 0 || padLen > blockSize {
		return nil, errors.New("invalid padding")
	}
	for i := 0; i < padLen; i++ {
		if data[length-1-i] != byte(padLen) {
			return nil, errors.New("invalid padding")
		}
	}
	return data[:length-padLen], nil
}

// AES-256-CBC 加密
func Aes256CbcEncrypt(plainText, iv []byte) ([]byte, error) {
	if len(aes_key) != 32 {
		return nil, errors.New("key length must be 32 bytes")
	}
	if len(iv) != aes.BlockSize {
		return nil, errors.New("IV length must be 16 bytes")
	}

	block, err := aes.NewCipher(aes_key)
	if err != nil {
		return nil, err
	}

	// PKCS7 填充
	padded := PKCS7Pad(plainText, aes.BlockSize)

	cipherText := make([]byte, len(padded))
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(cipherText, padded)

	return cipherText, nil
}

// AES-256-CBC 解密
func Aes256CbcDecrypt(cipherText []byte, iv []byte) ([]byte, error) {
	if len(aes_key) != 32 {
		return nil, errors.New("key length must be 32 bytes")
	}
	if len(iv) != aes.BlockSize {
		return nil, errors.New("IV length must be 16 bytes")
	}

	block, err := aes.NewCipher(aes_key)
	if err != nil {
		return nil, err
	}

	if len(cipherText)%aes.BlockSize != 0 {
		return nil, fmt.Errorf("cipherText is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(cipherText))
	mode.CryptBlocks(plaintext, cipherText)

	// 去掉 PKCS7 填充
	plaintext, err = PKCS7Unpad(plaintext, aes.BlockSize)
	if err != nil {
		return nil, err
	}

	return plaintext, nil
}

// 加密数据
func (x XlogEncrypt) EncryptData(data []byte) ([]byte, error) {
	encrypted, err := Aes256CbcEncrypt(data, x.AesIv)
	if err != nil {
		return nil, err
	}

	return encrypted, nil
}

// 解密数据
func (x XlogEncrypt) DecryptData(data []byte) ([]byte, error) {
	decrypted, err := Aes256CbcDecrypt(data, x.AesIv)
	if err != nil {
		return nil, err
	}

	return decrypted, nil
}