package crypto

import (
	"context"
	"fmt"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
)

// PostgresKeyStore reads/writes the crypto_keys table that records which key
// versions exist and which is currently active (the write key). The key
// material is NOT stored here — only version metadata.
type PostgresKeyStore struct {
	pool poolExec
}

// poolExec is the subset of *pgxpool.Pool used by the store. Declared as an
// interface so tests can inject a fake without a live DB.
type poolExec interface {
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

// poolFullExec additionally allows mutations (used by RotateActive/RetireVersion).
// Exec returns the CONCRETE pgconn.CommandTag so *pgxpool.Pool satisfies this
// interface — an interface-typed return would not match pgxpool's method
// signature, making the type assertion in RotateActive/RetireVersion silently
// fail and key rotation always error out.
type poolFullExec interface {
	poolExec
	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
}

// NewPostgresKeyStore builds the store over a pgx pool (or any compatible
// querier). It satisfies KeyStore.
func NewPostgresKeyStore(pool poolExec) *PostgresKeyStore {
	return &PostgresKeyStore{pool: pool}
}

// ActiveVersion returns the single crypto_keys row whose status='active'. If
// none is found it falls back to version 1 (single-key bootstrap before the
// seed row exists).
func (s *PostgresKeyStore) ActiveVersion(ctx context.Context) (int, error) {
	var v int
	err := s.pool.QueryRow(ctx,
		`SELECT version FROM crypto_keys WHERE status = 'active' ORDER BY version DESC LIMIT 1`,
	).Scan(&v)
	if err != nil {
		if err == pgx.ErrNoRows {
			return 1, nil
		}
		return 0, fmt.Errorf("crypto key active version: %w", err)
	}
	return v, nil
}

// RotateActive bumps the active write version: it retires the current active
// row and inserts/activates the next version (current+1). The new version's key
// material must already be loadable by the provider (e.g. APP_MASTER_KEY_V<n>
// staged) before this is called. Returns the new active version.
//
// pool must support Exec (a *pgxpool.Pool does). Both rows are written in one
// transaction by the caller's pool if it is transactional; here we issue two
// statements and rely on the unique active invariant being eventually consistent
// within the rotation window (both versions remain loadable, so decrypts never
// break).
func (s *PostgresKeyStore) RotateActive(ctx context.Context, provider string, keyRef string) (int, error) {
	fe, ok := s.pool.(poolFullExec)
	if !ok {
		return 0, fmt.Errorf("crypto key store: pool does not support writes")
	}
	cur, err := s.ActiveVersion(ctx)
	if err != nil {
		return 0, err
	}
	next := cur + 1
	if _, err := fe.Exec(ctx,
		`UPDATE crypto_keys SET status = 'retiring' WHERE status = 'active'`); err != nil {
		return 0, fmt.Errorf("retire current key: %w", err)
	}
	if _, err := fe.Exec(ctx,
		`INSERT INTO crypto_keys (version, provider, key_ref, status)
		 VALUES ($1, $2, NULLIF($3, ''), 'active')
		 ON CONFLICT (version) DO UPDATE SET status = 'active', provider = EXCLUDED.provider, key_ref = EXCLUDED.key_ref`,
		next, provider, keyRef); err != nil {
		return 0, fmt.Errorf("insert next key: %w", err)
	}
	return next, nil
}

// RetireVersion marks a fully-superseded version as retired (no longer needed
// after the re-encrypt job has moved every row off it). The provider may then
// safely drop that key material.
func (s *PostgresKeyStore) RetireVersion(ctx context.Context, version int) error {
	fe, ok := s.pool.(poolFullExec)
	if !ok {
		return fmt.Errorf("crypto key store: pool does not support writes")
	}
	if _, err := fe.Exec(ctx,
		`UPDATE crypto_keys SET status = 'retired' WHERE version = $1`, version); err != nil {
		return fmt.Errorf("retire key v%d: %w", version, err)
	}
	return nil
}

var _ KeyStore = (*PostgresKeyStore)(nil)