package crypto

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

// BackfillEncryptedKeys encrypts any remaining plaintext API keys in
// provider_configs.api_key_plain using the env-derived key, writing the
// result to api_key_enc (BYTEA) + key_nonce. Clears api_key_plain after.
//
// This runs once on upgrade when ENCRYPTION_KEY is first set. Subsequent starts are
// a no-op (api_key_plain will be NULL for all rows).
//
// Personal-scope keys are temporarily encrypted with the env key. They get
// re-encrypted with the user's UEK on next login (see VaultUpgradeOnLogin).
func BackfillEncryptedKeys(db *sql.DB, encryptionKey string) error {
	envKey, err := DeriveKeyFromEnv(encryptionKey)
	if err != nil {
		return fmt.Errorf("vault backfill: %w", err)
	}

	// Check if api_key_plain column exists (it won't on fresh installs after
	// the column is eventually dropped)
	var colExists bool
	db.QueryRow(`
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_name = 'provider_configs' AND column_name = 'api_key_plain'
		)
	`).Scan(&colExists)
	if !colExists {
		return nil // Column already dropped — nothing to backfill
	}

	// Find rows with plaintext keys that haven't been encrypted yet
	rows, err := db.Query(`
		SELECT id, api_key_plain
		FROM provider_configs
		WHERE api_key_plain IS NOT NULL
		  AND api_key_plain != ''
		  AND api_key_enc IS NULL
	`)
	if err != nil {
		return fmt.Errorf("vault backfill query: %w", err)
	}
	defer rows.Close()

	var count int
	for rows.Next() {
		var id, plaintext string
		if err := rows.Scan(&id, &plaintext); err != nil {
			return fmt.Errorf("vault backfill scan: %w", err)
		}

		ciphertext, nonce, err := Encrypt(plaintext, envKey)
		if err != nil {
			return fmt.Errorf("vault backfill encrypt %s: %w", id, err)
		}

		_, err = db.Exec(`
			UPDATE provider_configs
			SET api_key_enc = $1, key_nonce = $2, api_key_plain = NULL
			WHERE id = $3
		`, ciphertext, nonce, id)
		if err != nil {
			return fmt.Errorf("vault backfill update %s: %w", id, err)
		}
		count++
	}

	if count > 0 {
		log.Printf("  🔐 Vault backfill: encrypted %d API key(s)", count)
	}
	return nil
}

// EnforceEncryptionKey checks startup preconditions:
//   - If encrypted keys exist in DB and ENCRYPTION_KEY is empty → fatal.
//   - If no keys exist and ENCRYPTION_KEY is empty → warn (fresh install OK).
func EnforceEncryptionKey(db *sql.DB, encryptionKey string) error {
	if db == nil {
		return nil // No DB, no keys to protect
	}

	// Check if any encrypted keys exist
	var hasEncrypted bool
	db.QueryRow(`
		SELECT EXISTS (
			SELECT 1 FROM provider_configs WHERE api_key_enc IS NOT NULL
		)
	`).Scan(&hasEncrypted)

	// Check if any plaintext keys exist (pre-backfill)
	var hasPlaintext bool
	// api_key_plain column may not exist on fresh installs
	var colExists bool
	db.QueryRow(`
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_name = 'provider_configs' AND column_name = 'api_key_plain'
		)
	`).Scan(&colExists)
	if colExists {
		db.QueryRow(`
			SELECT EXISTS (
				SELECT 1 FROM provider_configs
				WHERE api_key_plain IS NOT NULL AND api_key_plain != ''
			)
		`).Scan(&hasPlaintext)
	}

	if encryptionKey == "" {
		if hasEncrypted {
			return fmt.Errorf(
				"ENCRYPTION_KEY is not set but encrypted API keys exist in the database. " +
					"Set ENCRYPTION_KEY to the value used when keys were encrypted, or data will be unrecoverable",
			)
		}
		if hasPlaintext {
			return fmt.Errorf(
				"ENCRYPTION_KEY is not set but plaintext API keys need encryption. " +
					"Set ENCRYPTION_KEY before starting (required+)",
			)
		}
		log.Println("  ⚠ ENCRYPTION_KEY not set — API key encryption disabled (OK for fresh install with no providers)")
	}

	return nil
}