package aws

import (
	"time"

	"github.com/jellydator/ttlcache/v2"
	"github.com/tidbcloud/pkgv2/secret/encrypt"
)

var (
	cacheSize = 10240
	cacheTTL  = 5 * time.Minute
)

// KmsEncryptor encrypts with AWS KMS
type KmsEncryptor struct {
	keyID    KeyID
	ttlCache *ttlcache.Cache
}

// Encrypt satisfies the Encryptor interface
func (aks KmsEncryptor) Encrypt(secret encrypt.Secret) ([]byte, error) {
	return EncryptDataWithKms(aks.keyID, []byte(secret))
}

// Decrypt satisfies the Encryptor interface
func (aks KmsEncryptor) Decrypt(cyphertext []byte) (encrypt.Secret, error) {
	// FAST PATH:
	secret, err := aks.ttlCache.Get(string(cyphertext))
	if err == nil {
		// get cleartext from cache successfully, return it
		// immediately instead of accessing KMS
		return secret.(encrypt.Secret), nil
	}

	// SLOW PATH:
	cleartext, err := DecryptDataWithKms(cyphertext)
	if err != nil {
		return "", err
	}

	res := encrypt.Secret(cleartext)
	// add the decryption result to athe cache to speed up the subsequent decryption
	_ = aks.ttlCache.Set(string(cyphertext), res)

	return res, err
}

// NewKmsEncryptor generates an KmsEncryptor
func NewKmsEncryptor(CMKArn string) (KmsEncryptor, error) {
	ttlCache := ttlcache.NewCache()
	ttlCache.SetCacheSizeLimit(cacheSize)
	_ = ttlCache.SetTTL(cacheTTL)

	keyID, err := GetKeyID(CMKArn)
	return KmsEncryptor{keyID: keyID, ttlCache: ttlCache}, err
}

var _ encrypt.Encryptor = KmsEncryptor{}