package envelope

import (
	"context"
	"crypto/rand"
	"fmt"

	"atlas9.dev/c/core"
)

// Provision generates, wraps, and stores a fresh DEK for a tenant. Call once
// when the tenant is created so an encryptor is available at read time; DEKs
// are never created lazily on the read/write path.
func Provision(ctx context.Context, deks DekStore, wrapper Wrapper, tenant core.ID) error {

	var raw [32]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return err
	}
	wrapped, err := wrapper.Wrap(ctx, raw[:])
	if err != nil {
		return fmt.Errorf("wrapping new dek: %w", err)
	}
	if _, err := deks.Create(ctx, tenant, wrapped); err != nil {
		return fmt.Errorf("storing new dek: %w", err)
	}
	return nil
}

// EncryptorFactory builds tenant-scoped encryptors: For loads and unwraps a
// tenant's DEK and returns an encryptor bound to it. It holds only the
// long-lived [Wrapper]; the per-transaction [DekStore] is passed to For. DEKs
// are created separately by [Provision] at tenant-creation time.
type EncryptorFactory struct {
	wrapper Wrapper
}

func NewEncryptorFactory(wrapper Wrapper) *EncryptorFactory {
	if wrapper == nil {
		panic("NewEncryptorFactory: wrapper must not be nil")
	}
	return &EncryptorFactory{wrapper: wrapper}
}

// For loads and unwraps the tenant's active DEK from deks, returning an
// encryptor bound to it. The plaintext DEK lives for the lifetime of the
// returned encryptor, so construct one per operation (e.g. per request) and
// reuse it across that operation's values — a List then unwraps the DEK once.
// Returns core.ErrNotFound if the tenant has no DEK (Provision was never
// called).
func (f *EncryptorFactory) For(ctx context.Context, deks DekStore, tenant core.ID) (*TenantEncryptor, error) {

	dek, err := deks.ForTenant(ctx, tenant)
	if err != nil {
		return nil, err
	}
	key, err := f.wrapper.Unwrap(ctx, dek.WrappedKey)
	if err != nil {
		return nil, fmt.Errorf("unwrapping dek: %w", err)
	}
	return &TenantEncryptor{dekID: dek.ID, plainKey: key}, nil
}

// TenantEncryptor seals and opens values for one tenant using that tenant's
// DEK, which [EncryptorFactory.For] unwraps once when building the encryptor
// and which is then held in memory for the encryptor's lifetime. Seal returns
// the DEK id and ciphertext to store as columns; Open takes them back.
type TenantEncryptor struct {
	dekID    core.ID
	plainKey []byte
}

// Seal encrypts plaintext under the tenant's DEK, returning the DEK id and the
// ciphertext (carrying the sealed-data header). Store both.
func (e *TenantEncryptor) Seal(plaintext []byte) (core.ID, []byte, error) {

	gcm, err := aesGCMSeal(e.plainKey, plaintext)
	if err != nil {
		return core.ID{}, nil, err
	}

	// framing: [magic][version][nonce || GCM ciphertext]
	ct := make([]byte, 0, headerLen+len(gcm))
	ct = append(ct, magic[:]...)
	ct = append(ct, version)
	ct = append(ct, gcm...)
	return e.dekID, ct, nil
}

// Open reverses Seal, decrypting with the DEK this encryptor holds — the one
// unwrapped when it was built (see [EncryptorFactory.For]). Open does no DEK
// lookup or unwrapping itself. Ciphertext sealed under a different DEK (e.g. a
// pre-rotation one) fails the GCM authentication check.
func (e *TenantEncryptor) Open(ciphertext []byte) ([]byte, error) {

	if err := validHeader(ciphertext); err != nil {
		return nil, fmt.Errorf("open: %w", err)
	}
	return aesGCMOpen(e.plainKey, ciphertext[headerLen:])
}