package ole

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"

	"code.justin.tv/amzn/TwitchOLE/ole/internal/datakeycache"
	"code.justin.tv/amzn/TwitchOLE/ole/internal/encryptedobject"
	"code.justin.tv/amzn/TwitchOLE/ole/internal/encryptioncontext"
	"code.justin.tv/amzn/TwitchOLE/ole/internal/keybytes"
	"github.com/golang/protobuf/proto"
)

type decryptionReader struct {
	k    *KMSClient
	r    io.Reader
	read bool
	buf  *bytes.Reader
}

// Read is not safe to read concurrently
func (dr *decryptionReader) Read(p []byte) (n int, err error) {
	if !dr.read {
		bs, err := ioutil.ReadAll(dr.r)
		dr.read = true
		if err != nil {
			return 0, err
		}

		obj := &encryptedobject.EncryptedObject{}
		err = proto.Unmarshal(bs, obj)
		if err != nil {
			return 0, err
		}

		var plaintext []byte
		plaintext, err = dr.k.decrypt(obj)
		if err != nil {
			return 0, err
		}

		dr.buf = bytes.NewReader(plaintext)
	}

	if dr.buf == nil {
		return 0, io.EOF
	}

	return dr.buf.Read(p)
}

func (k *KMSClient) decrypt(obj *encryptedobject.EncryptedObject) (plaintext []byte, err error) {
	awsKMSEncryptedDataKey := obj.GetAWSKMSEncryptedDataKey()
	if awsKMSEncryptedDataKey == nil {
		return nil, errors.New("encrypted object is missing encrypted data key field")
	}

	dataKey, err := k.getOrDecryptDecryptionDataKey(datakeycache.DecryptionKeyCacheCompositeKey{
		EncryptedDataKey:  awsKMSEncryptedDataKey.EncryptedDataKey,
		KeyProviderID:     "aws-kms",
		KeyInfo:           awsKMSEncryptedDataKey.CmkArn,
		AlgorithmID:       awsKMSEncryptedDataKey.AlgorithmId.String(),
		EncryptionContext: awsKMSEncryptedDataKey.EncryptionContext,
	})
	if err != nil {
		return nil, err
	}
	defer keybytes.Zero(dataKey)

	ec := encryptioncontext.Serialize(awsKMSEncryptedDataKey.EncryptionContext)
	return k.cfg.Algorithm.Decrypt(dataKey, obj.GetCiphertext(), obj.GetIv(), ec)
}
