package main

import (
	"log"
	"math/rand"
	"time"

	"code.justin.tv/eventbus/client/encryption"
	"code.justin.tv/eventbus/schema/pkg/eventbus/authorization"

	"github.com/Pallinder/go-randomdata"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
)

const (
	CMKArn    = "arn:aws:kms:us-west-2:297385687169:key/5d00e4b7-4621-481a-8222-f188be05f938"
	AWSRegion = "us-west-2"
)

type SingleTrialReport struct {
	BytesEncrypted uint64
	BytesDecrypted uint64
	EncryptCalls   uint64
	DecryptCalls   uint64
	TimeElapsed    time.Duration
}

type OLETester struct {
	client             *encryption.AuthorizedFieldClient
	encryptionContexts []authorization.Context
}

func NewOLETester(config *LoadTestConfig) *OLETester {
	session := session.Must(session.NewSession(aws.NewConfig().WithRegion(AWSRegion)))
	client, err := encryption.NewAuthorizedFieldClient(session, encryption.WithCMKArn(CMKArn))
	if err != nil {
		log.Fatal("failed to create authorized field client", err)
	}

	encryptionContexts := make([]authorization.Context, config.EncryptionContexts)
	for i := range encryptionContexts {
		encryptionContexts[i] = authorization.Context{
			authorization.EventType:   randomdata.SillyName(),
			authorization.Environment: randomdata.SillyName(),
			authorization.MessageName: randomdata.SillyName(),
			authorization.FieldName:   randomdata.SillyName(),
		}
	}

	return &OLETester{
		client:             client,
		encryptionContexts: encryptionContexts,
	}
}

func (o *OLETester) E2E(testType string) *SingleTrialReport {
	plaintext := randomdata.Paragraph()
	result := &SingleTrialReport{}
	if testType == "encryptdecrypt" {
		start := time.Now()
		encrypted := o.Encrypt(plaintext)
		decrypted := o.Decrypt(encrypted)
		if plaintext != decrypted {
			log.Printf("expected %q, but got %q\n", plaintext, decrypted)
		}
		result.BytesEncrypted = uint64(len([]byte(plaintext)))
		result.BytesDecrypted = uint64(len(encrypted.OlePayload))
		result.TimeElapsed = time.Since(start)
		result.EncryptCalls = 1
		result.DecryptCalls = 1
	}
	return result
}

func (o *OLETester) Encrypt(plaintext string) *authorization.String {
	encCtx := o.encryptionContext()
	encrypted, err := o.client.EncryptString(encCtx, plaintext)
	if err != nil {
		log.Printf("failed to encrypt\n\terror: %s\n\tcontext: %+v\n\ttext: %s\n", err, encCtx, plaintext)
	}
	return encrypted
}

func (o *OLETester) Decrypt(blob *authorization.String) string {
	decrypted, err := o.client.DecryptString(blob)
	if err != nil {
		log.Printf("failed to decrypt\n\terror: %s\n\tblob: %s\n", err, string(blob.OlePayload))
	}

	return decrypted
}

func (o *OLETester) encryptionContext() authorization.Context {
	return o.encryptionContexts[rand.Intn(len(o.encryptionContexts))]
}
