package secure

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"time"

	"provenance/internal/storage"
)

const KeyBytes = 32

type Vault struct {
	store storage.Secrets
	aead  cipher.AEAD
}

func New(store storage.Secrets, key string) (*Vault, error) {
	raw, err := base64.StdEncoding.DecodeString(key)
	if err != nil || len(raw) != KeyBytes {
		return nil, errors.New("encryption key must contain 32 base64-encoded bytes")
	}
	block, err := aes.NewCipher(raw)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &Vault{store, aead}, nil
}

func Token() string { return rand.Text() }

func (v *Vault) Put(ctx context.Context, kind, id string, data []byte, expires time.Time) error {
	sealed, err := v.seal(kind, id, data)
	if err != nil {
		return err
	}
	return v.store.Put(ctx, kind, id, sealed, expires)
}

func (v *Vault) Add(ctx context.Context, kind, id string, data []byte, expires time.Time) error {
	sealed, err := v.seal(kind, id, data)
	if err != nil {
		return err
	}
	return v.store.Add(ctx, kind, id, sealed, expires)
}

func (v *Vault) Get(ctx context.Context, kind, id string) ([]byte, error) {
	data, err := v.store.Get(ctx, kind, id)
	if err != nil {
		return nil, err
	}
	return v.open(kind, id, data)
}

func (v *Vault) Consume(ctx context.Context, kind, id string) ([]byte, error) {
	data, err := v.store.Consume(ctx, kind, id)
	if err != nil {
		return nil, err
	}
	return v.open(kind, id, data)
}

func (v *Vault) Remove(ctx context.Context, kind, id string) error {
	return v.store.Remove(ctx, kind, id)
}

func (v *Vault) seal(kind, id string, data []byte) ([]byte, error) {
	nonce := make([]byte, v.aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}
	return v.aead.Seal(nonce, nonce, data, []byte(kind+"\x00"+id)), nil
}

func (v *Vault) open(kind, id string, data []byte) ([]byte, error) {
	size := v.aead.NonceSize()
	if len(data) < size {
		return nil, errors.New("invalid encrypted data")
	}
	return v.aead.Open(nil, data[:size], data[size:], []byte(kind+"\x00"+id))
}