package jwe

import (
	"context"
	"sync"

	"github.com/lestrrat-go/jwx/internal/base64"
	"github.com/lestrrat-go/jwx/jwa"
	"github.com/pkg/errors"
)

var encryptCtxPool = sync.Pool{
	New: func() interface{} {
		return &encryptCtx{}
	},
}

func getEncryptCtx() *encryptCtx {
	return encryptCtxPool.Get().(*encryptCtx)
}

func releaseEncryptCtx(ctx *encryptCtx) {
	ctx.protected = nil
	ctx.contentEncrypter = nil
	ctx.generator = nil
	ctx.keyEncrypters = nil
	ctx.compress = jwa.NoCompress
	encryptCtxPool.Put(ctx)
}

// Encrypt takes the plaintext and encrypts into a JWE message.
func (e encryptCtx) Encrypt(plaintext []byte) (*Message, error) {
	bk, err := e.generator.Generate()
	if err != nil {
		return nil, errors.Wrap(err, "failed to generate key")
	}
	cek := bk.Bytes()

	if e.protected == nil {
		// shouldn't happen, but...
		e.protected = NewHeaders()
	}

	if err := e.protected.Set(ContentEncryptionKey, e.contentEncrypter.Algorithm()); err != nil {
		return nil, errors.Wrap(err, `failed to set "enc" in protected header`)
	}

	compression := e.compress
	if compression != jwa.NoCompress {
		if err := e.protected.Set(CompressionKey, compression); err != nil {
			return nil, errors.Wrap(err, `failed to set "zip" in protected header`)
		}
	}

	// In JWE, multiple recipients may exist -- they receive an
	// encrypted version of the CEK, using their key encryption
	// algorithm of choice.
	recipients := make([]Recipient, len(e.keyEncrypters))
	for i, enc := range e.keyEncrypters {
		r := NewRecipient()
		if err := r.Headers().Set(AlgorithmKey, enc.Algorithm()); err != nil {
			return nil, errors.Wrap(err, "failed to set header")
		}
		if v := enc.KeyID(); v != "" {
			if err := r.Headers().Set(KeyIDKey, v); err != nil {
				return nil, errors.Wrap(err, "failed to set header")
			}
		}

		enckey, err := enc.Encrypt(cek)
		if err != nil {
			return nil, errors.Wrap(err, `failed to encrypt key`)
		}
		if enc.Algorithm() == jwa.ECDH_ES || enc.Algorithm() == jwa.DIRECT {
			if len(e.keyEncrypters) > 1 {
				return nil, errors.Errorf("unable to support multiple recipients for ECDH-ES")
			}
			cek = enckey.Bytes()
		} else {
			if err := r.SetEncryptedKey(enckey.Bytes()); err != nil {
				return nil, errors.Wrap(err, "failed to set encrypted key")
			}
		}
		if hp, ok := enckey.(populater); ok {
			if err := hp.Populate(r.Headers()); err != nil {
				return nil, errors.Wrap(err, "failed to populate")
			}
		}
		recipients[i] = r
	}

	// If there's only one recipient, you want to include that in the
	// protected header
	if len(recipients) == 1 {
		h, err := e.protected.Merge(context.TODO(), recipients[0].Headers())
		if err != nil {
			return nil, errors.Wrap(err, "failed to merge protected headers")
		}
		e.protected = h
	}

	aad, err := e.protected.Encode()
	if err != nil {
		return nil, errors.Wrap(err, "failed to base64 encode protected headers")
	}

	plaintext, err = compress(plaintext, compression)
	if err != nil {
		return nil, errors.Wrap(err, `failed to compress payload before encryption`)
	}

	// ...on the other hand, there's only one content cipher.
	iv, ciphertext, tag, err := e.contentEncrypter.Encrypt(cek, plaintext, aad)
	if err != nil {
		return nil, errors.Wrap(err, "failed to encrypt payload")
	}

	msg := NewMessage()

	decodedAad, err := base64.Decode(aad)
	if err != nil {
		return nil, errors.Wrap(err, "failed to decode base64")
	}
	if err := msg.Set(AuthenticatedDataKey, decodedAad); err != nil {
		return nil, errors.Wrapf(err, `failed to set %s`, AuthenticatedDataKey)
	}
	if err := msg.Set(CipherTextKey, ciphertext); err != nil {
		return nil, errors.Wrapf(err, `failed to set %s`, CipherTextKey)
	}
	if err := msg.Set(InitializationVectorKey, iv); err != nil {
		return nil, errors.Wrapf(err, `failed to set %s`, InitializationVectorKey)
	}
	if err := msg.Set(ProtectedHeadersKey, e.protected); err != nil {
		return nil, errors.Wrapf(err, `failed to set %s`, ProtectedHeadersKey)
	}
	if err := msg.Set(RecipientsKey, recipients); err != nil {
		return nil, errors.Wrapf(err, `failed to set %s`, RecipientsKey)
	}
	if err := msg.Set(TagKey, tag); err != nil {
		return nil, errors.Wrapf(err, `failed to set %s`, TagKey)
	}

	return msg, nil
}