package suite

import (
	"context"

	"code.justin.tv/eventbus/controlplane/internal/e2eaccounts"

	"code.justin.tv/eventbus/controlplane/e2e/internal/expected"

	"code.justin.tv/eventbus/controlplane/e2e/internal/e2eutil"
	"code.justin.tv/eventbus/controlplane/e2e/internal/report"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/pkg/errors"
)

var _ Runner = &ClientsTestSuite{}

// AuthorizedFieldGrantsTestSuite ensures that a couple created authorized field grants
// can be properly encrypted and decrypted. This test suite was written to run with
// the 09_authorized_field_grants seed data
type AuthorizedFieldGrantsTestSuite struct {
	*DefaultTestSuite // Use the default suite as ground work
}

func NewAuthorizedFieldGrantsTestSuite(suiteName string) (Runner, error) {
	// Base setup, testing, and cleaning procedures
	defaultTestSuite, err := NewDefaultTestSuite(suiteName)
	if err != nil {
		return nil, err
	}
	// Wrap the basic tooling with some extra tests
	return &AuthorizedFieldGrantsTestSuite{
		DefaultTestSuite: defaultTestSuite,
	}, nil
}

// Setup is nothing on top of the default setup phase
func (t *AuthorizedFieldGrantsTestSuite) Setup(ctx context.Context) report.Error {
	ctx = e2eutil.AppendTestPath(ctx, t.TestName())
	return t.DefaultTestSuite.Setup(ctx)
}

// Test includes the default testing, plus extra tests to publish
// and receive an event using the client libraries
func (t *AuthorizedFieldGrantsTestSuite) Test(ctx context.Context) {
	ctx = e2eutil.AppendTestPath(ctx, t.TestName())
	t.DefaultTestSuite.Test(ctx)

	t.Log(ctx, "Verifying encrypt and decrypt operations succeed and fail according to grants")

	pubCreds := e2eaccounts.AccountCredentials().PublisherConfig()
	subCreds := e2eaccounts.AccountCredentials().SubscriberConfig()

	// Make some encrypted payloads, one for each authorized field in the seed
	encryptRequests := []*struct {
		EventType   string
		Environment string
		MessageName string
		FieldName   string
		DataKey     *kms.GenerateDataKeyOutput
	}{
		{"E2ECreate", "production", "TopSecret", "AuthorizedEyesOnly", nil},
		{"E2EUpdate", "production", "TopSecret", "AuthorizedEyesOnly", nil},
		{"E2EUpdate", "production", "Foobar", "Garply", nil},
		{"E2EUpdate", "staging", "TopSecret", "AuthorizedEyesOnly", nil},
	}

	for _, req := range encryptRequests {
		eventType := e2eutil.Suffix(req.EventType, e2eutil.JobID(ctx))
		resp, err := encrypt(e2eaccounts.AccountCredentials().MainConfig(), expected.EncryptionContext(eventType, req.Environment, req.MessageName, req.FieldName))
		if err != nil {
			t.Error(ctx, "could not encrypt payloads during test setup", err)
			return
		}
		req.DataKey = resp
	}

	tt := []struct {
		Credentials   *aws.Config
		EventType     string
		Environment   string
		MessageName   string
		FieldName     string
		CanEncrypt    bool
		CanDecrypt    bool
		DecryptTarget *kms.GenerateDataKeyOutput // Use this to attempt a decrypt and check the result
	}{
		{pubCreds, "E2EUpdate", "production", "TopSecret", "AuthorizedEyesOnly", true, false, encryptRequests[1].DataKey},
		{pubCreds, "E2EUpdate", "production", "Foobar", "Garply", true, false, encryptRequests[2].DataKey},
		{pubCreds, "E2ECreate", "production", "TopSecret", "AuthorizedEyesOnly", false, false, encryptRequests[0].DataKey},
		{subCreds, "E2EUpdate", "production", "TopSecret", "AuthorizedEyesOnly", false, true, encryptRequests[1].DataKey},
		{subCreds, "E2EUpdate", "production", "Foobar", "Garply", false, false, encryptRequests[2].DataKey},
		{subCreds, "E2EUpdate", "staging", "TopSecret", "AuthorizedEyesOnly", false, false, encryptRequests[3].DataKey},
	}

	for _, e := range tt {
		eventType := e2eutil.Suffix(e.EventType, e2eutil.JobID(ctx))
		encryptionContext := expected.EncryptionContext(eventType, e.Environment, e.MessageName, e.FieldName)

		_, err := encrypt(e.Credentials, encryptionContext)
		if e.CanEncrypt && err != nil {
			t.Error(ctx, "expected no error on encrypt", err)
		} else if !e.CanEncrypt && err == nil {
			t.Error(ctx, "expected an error on encrypt", nil)
		}

		resp, err := decrypt(e.Credentials, encryptionContext, e.DecryptTarget.CiphertextBlob)
		if e.CanDecrypt && err != nil {
			t.Error(ctx, "expected no error on decrypt", err)
		} else if !e.CanDecrypt && err == nil {
			t.Error(ctx, "expected and error on decrypt", nil)
		} else if e.CanDecrypt {
			if string(resp.Plaintext) != string(e.DecryptTarget.Plaintext) {
				t.Error(ctx, "mismatch between expected and actual plaintext after decryption", nil)
			}
		}
	}
}

func (t *AuthorizedFieldGrantsTestSuite) Clean(ctx context.Context) []report.Error {
	ctx = e2eutil.AppendTestPath(ctx, t.TestName())
	return t.DefaultTestSuite.Clean(ctx)
}

func (t *AuthorizedFieldGrantsTestSuite) TestName() string {
	return "AuthorizedFieldGrantsTestSuite"
}

// will attempt Encrypt, GenerateDataKey and GenerateDataKeyWithoutPlaintext, returning the result of GenerateDataKey
func encrypt(c *aws.Config, encCtx map[string]string) (*kms.GenerateDataKeyOutput, error) {
	kmsClient := kms.New(session.Must(session.NewSession(c)))

	resp, err := kmsClient.GenerateDataKey(&kms.GenerateDataKeyInput{
		EncryptionContext: aws.StringMap(encCtx),
		KeyId:             aws.String(expected.AuthorizedFieldKeyARN),
		KeySpec:           aws.String("AES_256"),
	})
	if err != nil {
		return nil, errors.Wrap(err, "failed to call kms:GenerateDataKey")
	}

	// Also make sure GenerateDataKeyWithoutPlaintext work without error
	_, err = kmsClient.GenerateDataKeyWithoutPlaintext(&kms.GenerateDataKeyWithoutPlaintextInput{
		EncryptionContext: aws.StringMap(encCtx),
		KeyId:             aws.String(expected.AuthorizedFieldKeyARN),
		KeySpec:           aws.String("AES_256"),
	})
	if err != nil {
		return nil, errors.Wrap(err, "failed to call kms:GenerateDataKeyWithoutPlaintext")
	}

	_, err = kmsClient.Encrypt(&kms.EncryptInput{
		EncryptionContext: aws.StringMap(encCtx),
		KeyId:             aws.String(expected.AuthorizedFieldKeyARN),
		Plaintext:         []byte("some jargon to encrypt"),
	})
	if err != nil {
		return nil, errors.Wrap(err, "failed to call kms:Encrypt")
	}

	return resp, nil
}

// will attempt Decrypt, returning the result
func decrypt(c *aws.Config, encCtx map[string]string, blob []byte) (*kms.DecryptOutput, error) {
	if blob == nil {
		return nil, errors.New("received empty payload to decrypt")
	}

	kmsClient := kms.New(session.Must(session.NewSession(c)))

	resp, err := kmsClient.Decrypt(&kms.DecryptInput{
		EncryptionContext: aws.StringMap(encCtx),
		CiphertextBlob:    blob,
	})
	if err != nil {
		return nil, errors.Wrap(err, "failed to call kms:Decrypt")
	}

	return resp, err
}
