package cryptox

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

var ErrKeyUnavailable = errors.New("encryption key unavailable")

type AESGCM struct {
	aead    cipher.AEAD
	version int
	purpose string
}

func NewAESGCM(encodedKey string, version int, purpose string) (*AESGCM, error) {
	if encodedKey == "" {
		return nil, nil
	}
	key, err := base64.StdEncoding.DecodeString(encodedKey)
	if err != nil {
		key, err = base64.RawStdEncoding.DecodeString(encodedKey)
	}
	if err != nil || len(key) != 32 {
		return nil, errors.New("encryption key must be base64-encoded 32 bytes")
	}
	if version < 1 || purpose == "" {
		return nil, errors.New("encryption version and purpose are required")
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &AESGCM{aead: aead, version: version, purpose: purpose}, nil
}

func (c *AESGCM) Encrypt(plaintext []byte) ([]byte, int, error) {
	if c == nil {
		return nil, 0, ErrKeyUnavailable
	}
	nonce := make([]byte, c.aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, 0, err
	}
	ciphertext := c.aead.Seal(nil, nonce, plaintext, c.additionalData())
	return append(nonce, ciphertext...), c.version, nil
}

func (c *AESGCM) Decrypt(encrypted []byte, version int) ([]byte, error) {
	if c == nil {
		return nil, ErrKeyUnavailable
	}
	if version != c.version {
		return nil, fmt.Errorf("encryption key version %d is not loaded", version)
	}
	nonceSize := c.aead.NonceSize()
	if len(encrypted) <= nonceSize {
		return nil, errors.New("encrypted value is truncated")
	}
	plaintext, err := c.aead.Open(nil, encrypted[:nonceSize], encrypted[nonceSize:], c.additionalData())
	if err != nil {
		return nil, errors.New("encrypted value authentication failed")
	}
	return plaintext, nil
}

func (c *AESGCM) additionalData() []byte {
	return []byte(fmt.Sprintf("ai-gateway/%s/v%d", c.purpose, c.version))
}