package encryption

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

const (
	currentVersion byte = 1
	keySize             = 32
)

type Cipher struct {
	gcm cipher.AEAD
}

func New(masterKey []byte) (Cipher, error) {
	if len(masterKey) != keySize {
		return Cipher{}, fmt.Errorf("master key must be %d bytes", keySize)
	}

	block, err := aes.NewCipher(masterKey)
	if err != nil {
		return Cipher{}, err
	}

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

	return Cipher{gcm: gcm}, nil
}

func NewFromBase64(encodedKey string) (Cipher, error) {
	encodedKey = strings.TrimSpace(encodedKey)
	if encodedKey == "" {
		return Cipher{}, fmt.Errorf("MASTER_KEY_BASE64 is required")
	}

	key, err := base64.StdEncoding.DecodeString(encodedKey)
	if err != nil {
		return Cipher{}, fmt.Errorf("MASTER_KEY_BASE64 must be base64 encoded: %w", err)
	}

	return New(key)
}

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

	ciphertext := make([]byte, 0, 1+len(nonce)+len(plaintext)+c.gcm.Overhead())
	ciphertext = append(ciphertext, currentVersion)
	ciphertext = append(ciphertext, nonce...)
	ciphertext = c.gcm.Seal(ciphertext, nonce, plaintext, nil)
	return ciphertext, nil
}

func (c Cipher) Decrypt(ciphertext []byte) ([]byte, error) {
	nonceSize := c.gcm.NonceSize()
	if len(ciphertext) < 1+nonceSize+c.gcm.Overhead() {
		return nil, fmt.Errorf("ciphertext is too short")
	}
	if ciphertext[0] != currentVersion {
		return nil, fmt.Errorf("unsupported ciphertext version")
	}

	nonce := ciphertext[1 : 1+nonceSize]
	encrypted := ciphertext[1+nonceSize:]
	return c.gcm.Open(nil, nonce, encrypted, nil)
}