// Copyright © 2018 Banzai Cloud
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package awskms

import (
	"context"
	"strings"

	"emperror.dev/errors"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/kms"
	"github.com/bank-vaults/bank-vaults/pkg/kv"
)

const (
	// SseAES256 is an algorithm that can be used for Server-Side Encryption in AWS S3 buckets
	SseAES256 string = "AES256"
	// SseKMS is an algorithm that can be used for Server-Side Encryption in AWS S3 buckets
	SseKMS string = "aws:kms"
)

type awsKMS struct {
	ctx        context.Context
	store      kv.Service
	kmsService *kms.Client

	kmsID             string
	encryptionContext map[string]string
}

var _ kv.Service = &awsKMS{}

// NewWithConfig creates a new kv.Service encrypted by AWS KMS with and existing AWS Session
func NewWithConfig(ctx context.Context, config aws.Config, store kv.Service, kmsID string, encryptionContext map[string]string) (kv.Service, error) {
	if kmsID == "" {
		return nil, errors.Errorf("invalid kmsID specified: '%s'", kmsID)
	}

	return &awsKMS{
		ctx:               ctx,
		store:             store,
		kmsService:        kms.NewFromConfig(config),
		kmsID:             kmsID,
		encryptionContext: encryptionContext,
	}, nil
}

// New creates a new kv.Service encrypted by AWS KMS
func New(ctx context.Context, store kv.Service, region string, kmsID string, encryptionContext map[string]string) (kv.Service, error) {
	config, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, errors.WrapIf(err, "failed to load AWS config")
	}

	return NewWithConfig(ctx, config, store, kmsID, encryptionContext)
}

func (a *awsKMS) decrypt(cipherText []byte) ([]byte, error) {
	out, err := a.kmsService.Decrypt(a.ctx, &kms.DecryptInput{
		KeyId:             aws.String(a.kmsID),
		CiphertextBlob:    cipherText,
		EncryptionContext: a.encryptionContext,
		GrantTokens:       []string{},
	})
	if err != nil {
		return nil, errors.WrapIf(err, "failed to decrypt with KMS client")
	}

	return []byte(strings.TrimSpace(string(out.Plaintext))), nil
}

func (a *awsKMS) Get(ctx context.Context, key string) ([]byte, error) {
	cipherText, err := a.store.Get(ctx, key)
	if err != nil {
		return nil, errors.WrapIf(err, "failed to get data for KMS client")
	}

	return a.decrypt(cipherText)
}

func (a *awsKMS) encrypt(plainText []byte) ([]byte, error) {
	out, err := a.kmsService.Encrypt(a.ctx, &kms.EncryptInput{
		KeyId:             aws.String(a.kmsID),
		Plaintext:         plainText,
		EncryptionContext: a.encryptionContext,
		GrantTokens:       []string{},
	})
	if err != nil {
		return nil, errors.WrapIf(err, "failed to encrypt with KMS client")
	}

	return out.CiphertextBlob, nil
}

func (a *awsKMS) Set(ctx context.Context, key string, val []byte) error {
	cipherText, err := a.encrypt(val)
	if err != nil {
		return err
	}

	return a.store.Set(ctx, key, cipherText)
}
