package common

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

type Aes struct {
	key    []byte
	iv     []byte
	method string
}

// NewAes creates a new Aes instance with the given configuration.
func NewAes(key, iv, method string) *Aes {
	return &Aes{
		key:    []byte(key),
		iv:     []byte(iv),
		method: method,
	}
}

// Encrypt encrypts the data using AES encryption.
func (a *Aes) Encrypt(data string) (string, error) {
	block, err := aes.NewCipher(a.key)

	if err != nil {
		return "", err
	}
	encrypter := cipher.NewCBCEncrypter(block, a.iv)
	paddedData := pkcs7Pad([]byte(data), aes.BlockSize)
	encrypted := make([]byte, len(paddedData))
	encrypter.CryptBlocks(encrypted, paddedData)

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

// Decrypt decrypts the data using AES decryption.
func (a *Aes) Decrypt(data string) (string, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return "", err
	}

	decrypter := cipher.NewCBCDecrypter(block, a.iv)
	encrypted, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return "", err
	}

	decrypted := make([]byte, len(encrypted))
	decrypter.CryptBlocks(decrypted, encrypted)
	decrypted = pkcs7Unpad(decrypted)

	return string(decrypted), nil
}

// pkcs7Pad pads the given data according to PKCS7 padding standard.
func pkcs7Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padded := append(data, bytes.Repeat([]byte{byte(padding)}, padding)...)
	return padded
}

// pkcs7Unpad removes PKCS7 padding from the given data.
func pkcs7Unpad(data []byte) []byte {
	length := len(data)
	unpadding := int(data[length-1])
	return data[:length-unpadding]
}