package security

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

	"github.com/pkg/errors"
)

type IEncryptor interface {
	EncryptChunk([]byte) []byte
	DecryptChunk([]byte) ([]byte, error)
	EncryptStringToBase64(string) string
	DecryptBase64ToString(string) (string, error)
	Close()
}

type Encryptor struct {
	key           []byte
	iv            []byte
	gcm_encryptor cipher.AEAD
	block_auth    []byte
}

// NewEncryptor create new GCM - expects the key to be 32 bytes
func NewEncryptor(encryption_key []byte, auth []byte, nonce []byte) (IEncryptor, []byte, error) {
	cypher, err := aes.NewCipher(encryption_key)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}
	gcm, err := cipher.NewGCM(cypher)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}
	if nonce != nil && len(nonce) != gcm.NonceSize() {
		return nil, nil, errors.Errorf("nonce is incorrect size %v", gcm.NonceSize())
	}
	if nonce == nil {
		nonce = make([]byte, gcm.NonceSize())
		_, err = io.ReadFull(rand.Reader, nonce)
		if err != nil {
			return nil, nil, errors.WithStack(err)
		}
	}
	return &Encryptor{key: encryption_key, iv: nonce, gcm_encryptor: gcm, block_auth: auth}, nonce, nil
}

// EncryptChunk takes a chunk and ecnrypts it using GCM
func (s *Encryptor) EncryptChunk(chunk []byte) []byte {
	return s.gcm_encryptor.Seal(chunk[:0], s.iv, chunk, s.block_auth)
}

// DecryptChunk takes a chunk and decrypts it using GCM
func (s *Encryptor) DecryptChunk(chunk []byte) ([]byte, error) {
	data, err := s.gcm_encryptor.Open(chunk[:0], s.iv, chunk, s.block_auth)
	return data, errors.WithStack(err)
}

// EncryptStringToBase64 takes a string as input and encrypts it and encodes it for transportability
func (s *Encryptor) EncryptStringToBase64(input string) string {
	data := s.EncryptChunk([]byte(input))
	return base64.RawStdEncoding.EncodeToString(data)
}

// DecryptStringToBase64 takes a string as input and decodes it and decrypts it from transport
func (s *Encryptor) DecryptBase64ToString(input string) (string, error) {
	decoded, err := base64.RawStdEncoding.DecodeString(input)
	if err != nil {
		return "", errors.Wrapf(err, "request %s", input)
	}

	decrypted, err := s.DecryptChunk(decoded)
	if err != nil {
		return "", err
	}

	return string(decrypted), nil
}

func (s *Encryptor) Close() {
	s.gcm_encryptor = nil
	s.iv = nil
	s.key = nil
}