package crypto

import (
	"database/sql"
	"fmt"
	"log"
)

// Rekey re-encrypts all global/team API keys from oldKey to newKey.
// Personal BYOK keys are unaffected (they're keyed to the user's UEK,
// not the environment-derived key).
//
// Runs in a single transaction — all-or-nothing. On failure, no keys
// are modified.
//
// Usage:
//
//	switchboard vault rekey
//	  ENCRYPTION_KEY     = current/old key
//	  NEW_ENCRYPTION_KEY = new key to re-encrypt with
func Rekey(db *sql.DB, oldEnvKey, newEnvKey string) error {
	if oldEnvKey == "" {
		return fmt.Errorf("ENCRYPTION_KEY (current key) is required")
	}
	if newEnvKey == "" {
		return fmt.Errorf("NEW_ENCRYPTION_KEY is required")
	}
	if oldEnvKey == newEnvKey {
		return fmt.Errorf("ENCRYPTION_KEY and NEW_ENCRYPTION_KEY are identical — nothing to do")
	}

	oldKey, err := DeriveKeyFromEnv(oldEnvKey)
	if err != nil {
		return fmt.Errorf("failed to derive old key: %w", err)
	}

	newKey, err := DeriveKeyFromEnv(newEnvKey)
	if err != nil {
		return fmt.Errorf("failed to derive new key: %w", err)
	}

	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback() // no-op after commit

	// Find all global/team keys that are encrypted
	rows, err := tx.Query(`
		SELECT id, api_key_enc, key_nonce
		FROM provider_configs
		WHERE key_scope IN ('global', 'team')
		  AND api_key_enc IS NOT NULL
	`)
	if err != nil {
		return fmt.Errorf("failed to query encrypted keys: %w", err)
	}

	type rekeyRow struct {
		id         string
		ciphertext []byte
		nonce      []byte
	}

	var toRekey []rekeyRow
	for rows.Next() {
		var r rekeyRow
		if err := rows.Scan(&r.id, &r.ciphertext, &r.nonce); err != nil {
			rows.Close()
			return fmt.Errorf("failed to scan row: %w", err)
		}
		toRekey = append(toRekey, r)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return fmt.Errorf("row iteration error: %w", err)
	}

	if len(toRekey) == 0 {
		log.Println("No global/team encrypted keys found — nothing to rekey.")
		return nil
	}

	log.Printf("Rekeying %d provider config(s)...", len(toRekey))

	// Decrypt with old key, re-encrypt with new key, update in place
	for _, r := range toRekey {
		plaintext, err := Decrypt(r.ciphertext, r.nonce, oldKey)
		if err != nil {
			return fmt.Errorf("failed to decrypt config %s with old key: %w (is ENCRYPTION_KEY correct?)", r.id, err)
		}

		newCiphertext, newNonce, err := Encrypt(plaintext, newKey)
		if err != nil {
			return fmt.Errorf("failed to re-encrypt config %s: %w", r.id, err)
		}

		_, err = tx.Exec(`
			UPDATE provider_configs
			SET api_key_enc = $1, key_nonce = $2
			WHERE id = $3
		`, newCiphertext, newNonce, r.id)
		if err != nil {
			return fmt.Errorf("failed to update config %s: %w", r.id, err)
		}
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit rekey transaction: %w", err)
	}

	log.Printf("✅ Successfully rekeyed %d provider config(s).", len(toRekey))
	log.Println("Update ENCRYPTION_KEY to the new value and restart the server.")
	return nil
}