package aws

import (
	"context"
	"sync"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/kms"
	"github.com/pingcap/errors"
	"k8s.io/utils/pointer"
)

var (
	getKmsSvcOnce sync.Once
	gKmsSvc       *kms.Client
)

func getKmsSvc() *kms.Client {
	getKmsSvcOnce.Do(func() {
		awsConfig, err := awsconfig.LoadDefaultConfig(context.TODO())
		if err != nil {
			panic("AWS configuration error, " + err.Error())
		}
		gKmsSvc = kms.NewFromConfig(awsConfig)
	})
	return gKmsSvc
}

func WarmupKmsSvc() {
	_ = getKmsSvc()
}

// KeyID is for a customer master key
type KeyID string

// GetKeyID get the key id
func GetKeyID(CMKArn string) (KeyID, error) {
	cmkID := KeyID(CMKArn)
	if len(cmkID) == 0 {
		return cmkID, errors.New("get cmk arn failed")
	}
	return cmkID, nil
}

// EncryptDataWithKms use kms to encrypt sensitive information
func EncryptDataWithKms(cmkID KeyID, plaintext []byte) ([]byte, error) {
	kmsSvc := getKmsSvc()

	// TODO: To improve security, we can specify the encryption context that will be used to encrypt the data.
	// When you use an encryption context to encrypt data, you must specify the same(an exact case-sensitive match)
	// encryption context to decrypt the data
	encryptInput := &kms.EncryptInput{
		KeyId:     pointer.String(string(cmkID)),
		Plaintext: plaintext,
	}

	result, err := kmsSvc.Encrypt(context.TODO(), encryptInput)
	if err != nil {
		return nil, errors.Annotate(err, "encrypt data with cmk id failed")
	}
	return result.CiphertextBlob, nil
}

// DecryptDataWithKms use kms to decrypt the encrypt the ciphertext
func DecryptDataWithKms(ciphertext []byte) ([]byte, error) {
	kmsSvc := getKmsSvc()

	decryptInput := &kms.DecryptInput{
		CiphertextBlob: ciphertext,
	}

	result, err := kmsSvc.Decrypt(context.TODO(), decryptInput)
	if err != nil {
		return nil, errors.Annotate(err, "decrypt ciphertext failed")
	}

	return result.Plaintext, nil
}