package crypto

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"sync"
)

// KeyStore tracks which key versions exist and which one is the current write
// ("active") version. The key material itself is never stored here — only the
// version metadata (see crypto_keys table). It is consulted on Encrypt to pick
// the active version and by the re-encrypt job to discover the target version.
type KeyStore interface {
	// ActiveVersion returns the version that new ciphertext must be sealed with.
	ActiveVersion(ctx context.Context) (int, error)
}

// staticKeyStore is a fixed-version store, used when no DB is wired (tests,
// degraded boot). It always reports the same active version.
type staticKeyStore struct{ v int }

// NewStaticKeyStore returns a KeyStore that always reports version v as active.
func NewStaticKeyStore(v int) KeyStore { return staticKeyStore{v: v} }

func (s staticKeyStore) ActiveVersion(context.Context) (int, error) { return s.v, nil }

// KeyedEncryptor seals/opens secrets with a versioned key. The on-disk format
// is unchanged from the legacy Encryptor — nonce(12) || ciphertext || tag —
// because the key version travels in a separate DB column (key_version), not in
// the ciphertext. This keeps every existing blob forward-compatible: a v1 blob
// is still openable as long as version 1 is recorded for its row.
type KeyedEncryptor struct {
	provider Provider
	store    KeyStore

	mu   sync.Mutex
	gcms map[int]cipher.AEAD // version -> AEAD, lazily built from provider keys
}

// NewKeyedEncryptor builds a KeyedEncryptor over a key provider and key store.
// store may be nil → a static version-1 store is used (back-compat with the
// single-key world).
func NewKeyedEncryptor(provider Provider, store KeyStore) *KeyedEncryptor {
	if store == nil {
		store = NewStaticKeyStore(1)
	}
	return &KeyedEncryptor{provider: provider, store: store, gcms: map[int]cipher.AEAD{}}
}

// gcmFor returns (and caches) the AEAD for a version, resolving the key bytes
// via the provider on first use.
func (e *KeyedEncryptor) gcmFor(ctx context.Context, version int) (cipher.AEAD, error) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if g, ok := e.gcms[version]; ok {
		return g, nil
	}
	key, err := e.provider.KeyMaterial(ctx, version)
	if err != nil {
		return nil, fmt.Errorf("resolve key v%d: %w", version, err)
	}
	if len(key) != 32 {
		return nil, fmt.Errorf("key v%d must be 32 bytes (got %d)", version, len(key))
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("aes new cipher v%d: %w", version, err)
	}
	g, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("new gcm v%d: %w", version, err)
	}
	e.gcms[version] = g
	return g, nil
}

// ActiveVersion reports the version new ciphertext is sealed with.
func (e *KeyedEncryptor) ActiveVersion(ctx context.Context) (int, error) {
	return e.store.ActiveVersion(ctx)
}

// EncryptVersion seals plaintext with the active version and returns the
// ciphertext together with the version it was sealed under (to persist into the
// row's key_version column).
func (e *KeyedEncryptor) EncryptVersion(ctx context.Context, plaintext []byte) ([]byte, int, error) {
	version, err := e.store.ActiveVersion(ctx)
	if err != nil {
		return nil, 0, fmt.Errorf("active key version: %w", err)
	}
	ct, err := e.sealWith(ctx, version, plaintext)
	if err != nil {
		return nil, 0, err
	}
	return ct, version, nil
}

// SealWithVersion encrypts plaintext under an explicit version. The re-encrypt
// job uses this to re-seal an old-version blob onto the new active version.
func (e *KeyedEncryptor) SealWithVersion(ctx context.Context, version int, plaintext []byte) ([]byte, error) {
	return e.sealWith(ctx, version, plaintext)
}

// sealWith encrypts with a specific version (used by EncryptVersion and the
// re-encrypt job which re-seals onto the new active version).
func (e *KeyedEncryptor) sealWith(ctx context.Context, version int, plaintext []byte) ([]byte, error) {
	g, err := e.gcmFor(ctx, version)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, g.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("rand nonce: %w", err)
	}
	return g.Seal(nonce, nonce, plaintext, nil), nil
}

// DecryptVersion opens a ciphertext that was sealed under the given version.
func (e *KeyedEncryptor) DecryptVersion(ctx context.Context, ciphertext []byte, version int) ([]byte, error) {
	g, err := e.gcmFor(ctx, version)
	if err != nil {
		return nil, err
	}
	if len(ciphertext) < g.NonceSize() {
		return nil, errors.New("ciphertext too short")
	}
	nonce, ct := ciphertext[:g.NonceSize()], ciphertext[g.NonceSize():]
	pt, err := g.Open(nil, nonce, ct, nil)
	if err != nil {
		return nil, fmt.Errorf("gcm open v%d: %w", version, err)
	}
	return pt, nil
}

// DecryptAny opens a ciphertext without knowing its key version: it tries the
// active version first, then every lower version down to 1, returning the first
// that authenticates. AES-GCM's auth tag guarantees only the correct key opens a
// blob, so trying versions is safe and lets the read path stay version-agnostic
// (callers need not thread key_version through every SELECT). Version count is
// tiny (1-2 in practice), so the loop is cheap. A version whose key material is
// not configured simply fails to build its AEAD and is skipped.
func (e *KeyedEncryptor) DecryptAny(ctx context.Context, ciphertext []byte) ([]byte, error) {
	active, err := e.store.ActiveVersion(ctx)
	if err != nil || active < 1 {
		active = 1
	}
	var lastErr error
	for v := active; v >= 1; v-- {
		pt, derr := e.DecryptVersion(ctx, ciphertext, v)
		if derr == nil {
			return pt, nil
		}
		lastErr = derr
	}
	if lastErr == nil {
		lastErr = errors.New("decrypt: no key version available")
	}
	return nil, lastErr
}