package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"io"
	"strings"

	"github.com/pkg/errors"
)

// make key into correct format
func aeadKey(key []byte) *[32]byte {
	var result [32]byte
	copy(result[:], key[:32])
	return &result
}

// Encrypt takes in a byte array (couchdb/redis data) + a key
// key must be 32 bytes long
// returns encrypted text, error
func Encrypt(data []byte, key []byte) (string, error) {
	ciphertext, err := encrypt(data, key)
	if err != nil {
		return "", errors.WithStack(err)
	}

	return base64.URLEncoding.EncodeToString(ciphertext), nil
}

// decrypts a given string of data + key, returned as a byte array
func Decrypt(data string, key []byte) ([]byte, error) {
	raw, err := base64.URLEncoding.DecodeString(data)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	plaintext, err := decrypt(raw, key)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return plaintext, nil
}

// encryptJson encrypts the json.RawMessage into a string then
// reformats the string encyrypted message back into the json format
func EncryptJSON(value json.RawMessage, key []byte) (json.RawMessage, error) {
	encrypted, err := Encrypt(value, key)
	if err != nil {
		return nil, err
	}

	encryptedJSON := `{ "EncryptedData": "` + encrypted + `" }`
	jsonMessage := json.RawMessage{}
	err = json.Unmarshal([]byte(encryptedJSON), &jsonMessage)
	if err != nil {
		return nil, err
	}

	return jsonMessage, nil
}

// decrypt json decrypts the encrypted data that was converted into a json.RawMessage and returns it
func DecryptJSON(value json.RawMessage, key []byte) (json.RawMessage, error) {
	msg := &map[string]string{}
	err := json.Unmarshal(value, msg)
	if err != nil {
		return nil, err
	}

	encryptedData := (*msg)["EncryptedData"]

	decryptedData, err := Decrypt(encryptedData, key)
	if err != nil {
		return nil, err
	}

	return decryptedData, nil
}

// EncryptRedis just adds the EncryptedData: to the front of the encrypted data
func EncryptRedis(value []byte, key []byte) (string, error) {
	data, err := Encrypt(value, key)
	if err != nil {
		return "", err
	}

	return "EncryptedData:" + data, nil
}

// decryptBtye removes the EncryptedData prefix, decrypts the data, then returns it
func DecryptRedis(value string, key []byte) ([]byte, error) {
	prefix := "EncryptedData:"

	if strings.HasPrefix(value, prefix) {
		// remove the prefix
		data := []byte(value[len(prefix):])

		// decrypt
		decryptedData, err := Decrypt(string(data), key)
		if err != nil {
			return nil, err
		}
		return decryptedData, nil
	}
	return nil, nil
}

// private function actually doing the encryption
func encrypt(plaintext []byte, key []byte) (ciphertext []byte, err error) {
	// key must be 32 bytes
	if len(key) != 32 {
		return nil, errors.Errorf("key must be exactly 32 bytes long, got %d bytes", len(key))
	}

	// transform key into *[32]byte
	encryptionKey := aeadKey(key)

	// create a new cipher block with 32 bit AES key
	block, err := aes.NewCipher(encryptionKey[:])
	if err != nil {
		return nil, errors.Errorf("Unable to create cipher block with key: %s", encryptionKey)
	}

	// returns the block cipher wrapped in Galois Counter Mode w/ standard nonce length.
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, errors.Errorf("Unable to create new GCM block")
	}

	nonce := make([]byte, gcm.NonceSize())

	// populate nonce with random data
	_, err = io.ReadFull(rand.Reader, nonce)
	if err != nil {
		return nil, errors.Errorf("Unable to ReadFull")
	}

	// encrypt and authenticate plaintext
	return gcm.Seal(nonce, nonce, plaintext, nil), nil
}

// private function actually doing the decryption
func decrypt(ciphertext []byte, key []byte) (plaintext []byte, err error) {
	// key must be 32 bytes
	if len(key) != 32 {
		return nil, errors.Errorf("key must be exactly 32 bytes long, got %d bytes", len(key))
	}

	// transform key into *[32]byte
	encryptionKey := aeadKey(key)

	// create a new cipher block with 32 bit AES key
	block, err := aes.NewCipher(encryptionKey[:])
	if err != nil {
		return nil, errors.Errorf("Unable to create cipher block with key: %s", encryptionKey)
	}

	// returns the block cipher wrapped in Galois/Counter Mode w/ standard nonce length.
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, errors.Errorf("Unable to create new GCM block")
	}

	if len(ciphertext) < gcm.NonceSize() {
		return nil, errors.Errorf("malformed ciphertext")
	}

	// decrypt and authenticate ciphertext; return updated slice
	return gcm.Open(nil,
		ciphertext[:gcm.NonceSize()],
		ciphertext[gcm.NonceSize():],
		nil,
	)
}