package aes

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

var _ Aes = (*aes)(nil)

type Aes interface {
	i()

	EncryptCBC(encryptStr string, urlEncode bool) (string, error)
	DecryptCBC(decryptStr string, urlEncode bool) (string, error)

	EncryptCFB(plain string, urlEncode bool) (string, error)
	DecryptCFB(encrypted string, urlEncode bool) (string, error)
}

type aes struct {
	key string
	iv  string
}

func New(key, iv string) Aes {
	return &aes{
		key: key,
		iv:  iv,
	}
}

func (a *aes) i() {}

func (a *aes) EncryptCBC(encryptStr string, urlEncode bool) (string, error) {
	encoder := base64.StdEncoding
	if urlEncode {
		encoder = base64.URLEncoding
	}

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

	blockSize := block.BlockSize()
	encryptBytes = pkcsPadding(encryptBytes, blockSize)
	blockMode := cipher.NewCBCEncrypter(block, []byte(a.iv))
	encrypted := make([]byte, len(encryptBytes))
	blockMode.CryptBlocks(encrypted, encryptBytes)

	return encoder.EncodeToString(encrypted), nil
}

func (a *aes) DecryptCBC(decryptStr string, urlEncode bool) (string, error) {
	encoder := base64.StdEncoding
	if urlEncode {
		encoder = base64.URLEncoding
	}

	decryptBytes, err := encoder.DecodeString(decryptStr)
	if err != nil {
		return "", err
	}

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

	blockMode := cipher.NewCBCDecrypter(block, []byte(a.iv))
	decrypted := make([]byte, len(decryptBytes))
	blockMode.CryptBlocks(decrypted, decryptBytes)
	decrypted = pkcsUnPadding(decrypted)

	return string(decrypted), nil
}

func (a *aes) EncryptCFB(plain string, urlEncode bool) (string, error) {
	encoder := base64.StdEncoding
	if urlEncode {
		encoder = base64.URLEncoding
	}

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

	encrypted := make([]byte, len(plain))
	stream := cipher.NewCFBEncrypter(block, []byte(a.iv))
	stream.XORKeyStream(encrypted, []byte(plain))

	return encoder.EncodeToString(encrypted), nil

}

func (a *aes) DecryptCFB(encrypted string, urlEncode bool) (string, error) {
	encoder := base64.StdEncoding
	if urlEncode {
		encoder = base64.URLEncoding
	}

	decryptBytes, err := encoder.DecodeString(encrypted)
	if err != nil {
		return "", err
	}

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

	plain := make([]byte, len(decryptBytes))
	stream := cipher.NewCFBDecrypter(block, []byte(a.iv))
	stream.XORKeyStream(plain, decryptBytes)

	return string(plain), nil
}