package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"fmt"

	"github.com/Luzifer/go-openssl/v3"
)

func addPKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padText...)
}

func removePKCS7Padding(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("data cannot be empty")
	}

	padding := int(data[len(data)-1])

	if padding > len(data) {
		return nil, fmt.Errorf("invalid padding size")
	}

	return data[:len(data)-padding], nil
}

// EncryptAES 加密函式
func EncryptAES(plaintext, key, iv []byte) ([]byte, error) {
	// 驗證密鑰長度是否符合 AES 的要求
	if len(key) != aes.BlockSize && len(key) != 16 && len(key) != 24 && len(key) != 32 {
		return nil, fmt.Errorf("invalid key size: %d, expected 16, 24, or 32", len(key))
	}

	// 創建 AES 密碼器
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
	}

	// 確保輸入的資料長度符合 block size，並進行填充
	blockSize := block.BlockSize()
	plaintext = addPKCS7Padding(plaintext, blockSize)

	// 創建 CBC 模式加密器
	blockMode := cipher.NewCBCEncrypter(block, iv)

	// 用來保存加密結果的緩衝區
	enc := make([]byte, len(plaintext))
	blockMode.CryptBlocks(enc, plaintext)

	return enc, nil
}

// DecryptAES 使用 AES CBC 模式進行解密
func DecryptAES(ciphertext, key, iv []byte) ([]byte, error) {
	// 驗證密鑰長度是否符合 AES 的要求
	if len(key) != aes.BlockSize && len(key) != 16 && len(key) != 24 && len(key) != 32 {
		return nil, fmt.Errorf("invalid key size: %d, expected 16, 24, or 32", len(key))
	}

	// 創建 AES 密碼器
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
	}

	// 確保輸入的資料長度符合 block size，並進行解密
	blockSize := block.BlockSize()
	blockMode := cipher.NewCBCDecrypter(block, iv[:blockSize])

	// 用來保存解密結果的緩衝區
	origData := make([]byte, len(ciphertext))
	blockMode.CryptBlocks(origData, ciphertext)

	// 去除 PKCS7 填充
	return removePKCS7Padding(origData)
}

// EncryptAESWithOpenSSL 使用Openssl 加密AES
func EncryptAESWithOpenSSL(value, key string) ([]byte, error) {
	o := openssl.New()

	enc, err := o.EncryptBytes(key, []byte(value), openssl.DigestMD5Sum)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt with OpenSSL: %w", err)
	}

	return enc, nil
}

// DecryptAESWithOpenSSL 使用Openssl 解密AES
func DecryptAESWithOpenSSL(value, key string) ([]byte, error) {
	o := openssl.New()

	dec, err := o.DecryptBytes(key, []byte(value), openssl.DigestMD5Sum)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt with OpenSSL: %w", err)
	}

	return dec, nil
}