package encryption

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"io"

	commonpb "go.temporal.io/api/common/v1"
	"go.temporal.io/sdk/converter"
)

const (
	encoding = "binary/encrypted"
)

func NewCodec(key []byte) *Codec {
	return &Codec{
		secretKey: key,
		metadata: map[string][]byte{
			converter.MetadataEncoding: []byte(encoding),
		},
	}
}

// Codec implements a Temporal PayloadCodec to encrypt the payloads.
// It uses AES crypt.
type Codec struct {
	secretKey []byte
	metadata  map[string][]byte
}

func (c *Codec) Encode(payloads []*commonpb.Payload) ([]*commonpb.Payload, error) {
	encryptedPayloads := make([]*commonpb.Payload, len(payloads))
	for i, p := range payloads {
		unencryptedBytes, err := p.Marshal()
		if err != nil {
			return nil, err
		}

		encryptedBytes, err := c.encrypt(unencryptedBytes)
		if err != nil {
			return nil, err
		}

		encryptedPayloads[i] = &commonpb.Payload{
			Metadata: c.metadata,
			Data:     encryptedBytes,
		}
	}

	return encryptedPayloads, nil
}

func (c *Codec) Decode(payloads []*commonpb.Payload) ([]*commonpb.Payload, error) {
	unencryptedPayloads := make([]*commonpb.Payload, len(payloads))
	for i, p := range payloads {
		if string(p.Metadata[converter.MetadataEncoding]) != encoding {
			unencryptedPayloads[i] = p
			continue
		}

		unencryptedBytes, err := c.decrypt(p.Data)
		if err != nil {
			return nil, err
		}

		unencryptedPayload := &commonpb.Payload{}
		err = unencryptedPayload.Unmarshal(unencryptedBytes)
		if err != nil {
			return nil, err
		}
		unencryptedPayloads[i] = unencryptedPayload
	}

	return unencryptedPayloads, nil
}

func (c *Codec) encrypt(b []byte) ([]byte, error) {
	block, err := aes.NewCipher(c.secretKey)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	return gcm.Seal(nonce, nonce, b, nil), nil
}

func (c *Codec) decrypt(b []byte) ([]byte, error) {
	block, err := aes.NewCipher(c.secretKey)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	if len(b) < nonceSize {
		return nil, errors.New("not enough encrypted data to decrypt")
	}

	nonce := b[:nonceSize]
	d := b[nonceSize:]

	return gcm.Open(nil, nonce, d, nil)
}
