package ole

import (
	"io"
	"time"

	"code.justin.tv/amzn/TwitchOLE/ole/internal/encryptedobject"
	"code.justin.tv/amzn/TwitchOLE/ole/internal/keybytes"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/google/uuid"
)

func (k *KMSClient) encrypt(plaintext []byte, encryptionContext map[string]string) (obj *encryptedobject.EncryptedObject, err error) {
	dataKey, encryptedDataKey, err := k.getOrGenerateEncryptionDataKey(encryptionContext)
	if err != nil {
		return nil, err
	}
	defer keybytes.Zero(dataKey)

	ciphertext, iv, err := k.cfg.Algorithm.Encrypt(dataKey, plaintext, encryptionContext)
	if err != nil {
		return nil, err
	}

	messageID, err := uuid.NewRandom()
	if err != nil {
		return nil, err
	}

	return &encryptedobject.EncryptedObject{
		Version:     1,
		MessageId:   []byte(messageID.String()),
		AlgorithmId: encryptedobject.Algorithm_AES_256_GCM,
		Timestamp:   &timestamp.Timestamp{Nanos: int32(time.Now().UnixNano())},
		Ciphertext:  ciphertext,
		Iv:          iv,
		DataKey: &encryptedobject.EncryptedObject_AWSKMSEncryptedDataKey{
			AWSKMSEncryptedDataKey: &encryptedobject.AWSKMSEncryptedDataKey{
				AlgorithmId:       encryptedobject.Algorithm_AES_256_GCM,
				EncryptionContext: encryptionContext,
				CmkArn:            k.cfg.CMKArn,
				EncryptedDataKey:  encryptedDataKey,
			},
		},
	}, nil
}

type encryptionWriter struct {
	k                 *KMSClient
	encryptionContext map[string]string
	w                 io.Writer
}

func (ew *encryptionWriter) Write(p []byte) (n int, err error) {
	obj, err := ew.k.encrypt(p, ew.encryptionContext)
	if err != nil {
		return 0, err
	}

	bs, err := proto.Marshal(obj)
	if err != nil {
		return 0, err
	}

	return ew.w.Write(bs)
}
