package menc

// 15% AI generated code
// mostly human made; AI comments and details

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"
	"errors"
)

var (
	ErrCipherTextTooShort = errors.New("ciphertext too short")
)

type AESGCM_AutoNonce struct {
	gcm      cipher.AEAD
	nonceSeq uint32
}

func NewAESGCM_AutoNonce(key []byte) (*AESGCM_AutoNonce, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &AESGCM_AutoNonce{
		gcm: gcm,
	}, nil
}

/*
Encrypt encrypts plaintext with optional AAD.

Nonce is sequencialy produced and auto managed.

Output format: nonce || ciphertext
*/
func (a *AESGCM_AutoNonce) Encrypt(plaintext, aad []byte) ([]byte, error) {
	ns := a.gcm.NonceSize()
	out := make([]byte, ns, ns+len(plaintext)+a.gcm.Overhead())
	nonce := out[:ns]

	binary.BigEndian.PutUint32(nonce[ns-4:], a.nonceSeq)
	a.nonceSeq++

	return a.gcm.Seal(out, nonce, plaintext, aad), nil
}

/*
Decrypt decrypts data produced by some AESGCM encrypt.

Expects input format: nonce || ciphertext
*/
func (a *AESGCM_AutoNonce) Decrypt(ciphertext, aad []byte) ([]byte, error) {
	ns := a.gcm.NonceSize()
	if len(ciphertext) < ns {
		return nil, ErrCipherTextTooShort
	}

	nonce := ciphertext[:ns]
	data := ciphertext[ns:]

	return a.gcm.Open(nil, nonce, data, aad)
}

/*
Encrypt encrypts plaintext with optional AAD.

A random nonce is geneated automaticaly.

Output format: nonce || ciphertext
*/
func AESGCM_Quick_Encrypt(key, plaintext, aad []byte) ([]byte, error) {
	gcm, err := helper_AESGCM_Quick(key)
	if err != nil {
		return nil, err
	}
	ns := gcm.NonceSize()
	out := make([]byte, ns, ns+len(plaintext)+gcm.Overhead())
	nonce := out[:ns]

	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}
	out = gcm.Seal(out, nonce, plaintext, aad)

	return out, nil
}

/*
Decrypt decrypts data produced by some AESGCM encrypt.

Expects input format: nonce || ciphertext
*/
func AESGCM_Quick_Decrypt(key, ciphertext, aad []byte) ([]byte, error) {
	gcm, err := helper_AESGCM_Quick(key)
	if err != nil {
		return nil, err
	}
	ns := gcm.NonceSize()
	if len(ciphertext) < ns {
		return nil, ErrCipherTextTooShort
	}

	nonce := ciphertext[:ns]
	data := ciphertext[ns:]

	return gcm.Open(nil, nonce, data, aad)
}

func helper_AESGCM_Quick(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}