package main

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/spf13/cobra"

	"git.wntrmute.dev/mc/mcdsl/terminal"
	"git.wntrmute.dev/mc/metacrypt/internal/config"
	"git.wntrmute.dev/mc/metacrypt/internal/crypto"
	"git.wntrmute.dev/mc/metacrypt/internal/db"
)

var migrateAADCmd = &cobra.Command{
	Use:   "migrate-aad",
	Short: "Migrate barrier entries to path-bound AAD encryption",
	Long: `Re-encrypts all barrier entries so that each entry's path is included
as GCM additional authenticated data (AAD). Entries already encrypted with
AAD are left untouched. Requires the unseal password.

This is a one-off migration for the nil-AAD to path-AAD transition. Run it
while the server is stopped to avoid concurrent access.`,
	RunE: runMigrateAAD,
}

var migrateAADDryRun bool

func init() {
	migrateAADCmd.Flags().BoolVar(&migrateAADDryRun, "dry-run", false, "report what would be migrated without writing")
	rootCmd.AddCommand(migrateAADCmd)
}

func runMigrateAAD(cmd *cobra.Command, args []string) error {
	configPath := cfgFile
	if configPath == "" {
		configPath = "/srv/metacrypt/metacrypt.toml"
	}

	cfg, err := config.Load(configPath)
	if err != nil {
		return err
	}

	database, err := db.Open(cfg.Database.Path)
	if err != nil {
		return err
	}
	defer func() { _ = database.Close() }()

	// Read unseal password.
	passwordBytes, err := terminal.ReadPasswordBytes("Unseal password: ")
	if err != nil {
		return fmt.Errorf("read password: %w", err)
	}
	defer crypto.Zeroize(passwordBytes)

	// Load seal config and derive MEK.
	mek, err := deriveMEK(database, passwordBytes)
	if err != nil {
		return err
	}
	defer crypto.Zeroize(mek)

	// Enumerate all barrier entries.
	ctx := context.Background()
	rows, err := database.QueryContext(ctx, "SELECT path, value FROM barrier_entries")
	if err != nil {
		return fmt.Errorf("query entries: %w", err)
	}
	defer func() { _ = rows.Close() }()

	type entry struct {
		path  string
		value []byte
	}
	var toMigrate []entry
	var alreadyMigrated, total int

	for rows.Next() {
		var path string
		var encrypted []byte
		if err := rows.Scan(&path, &encrypted); err != nil {
			return fmt.Errorf("scan entry: %w", err)
		}
		total++

		// Try decrypting with path AAD first (already migrated).
		if _, err := crypto.Decrypt(mek, encrypted, []byte(path)); err == nil {
			alreadyMigrated++
			continue
		}

		// Try decrypting with nil AAD (needs migration).
		if _, err := crypto.Decrypt(mek, encrypted, nil); err != nil {
			return fmt.Errorf("entry %q: cannot decrypt with AAD or without AAD — data may be corrupt", path)
		}

		toMigrate = append(toMigrate, entry{path: path, value: encrypted})
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("iterate entries: %w", err)
	}

	fmt.Printf("Total entries:      %d\n", total)
	fmt.Printf("Already migrated:   %d\n", alreadyMigrated)
	fmt.Printf("Need migration:     %d\n", len(toMigrate))

	if len(toMigrate) == 0 {
		fmt.Println("Nothing to migrate.")
		return nil
	}

	if migrateAADDryRun {
		fmt.Println("\nDry run — entries that would be migrated:")
		for _, e := range toMigrate {
			fmt.Printf("  %s\n", e.path)
		}
		return nil
	}

	// Migrate entries in a transaction.
	tx, err := database.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	stmt, err := tx.PrepareContext(ctx, "UPDATE barrier_entries SET value = ?, updated_at = datetime('now') WHERE path = ?")
	if err != nil {
		return fmt.Errorf("prepare update: %w", err)
	}
	defer func() { _ = stmt.Close() }()

	for _, e := range toMigrate {
		// Decrypt with nil AAD.
		plaintext, err := crypto.Decrypt(mek, e.value, nil)
		if err != nil {
			return fmt.Errorf("decrypt %q: %w", e.path, err)
		}

		// Re-encrypt with path AAD.
		newEncrypted, err := crypto.Encrypt(mek, plaintext, []byte(e.path))
		crypto.Zeroize(plaintext)
		if err != nil {
			return fmt.Errorf("encrypt %q: %w", e.path, err)
		}

		if _, err := stmt.ExecContext(ctx, newEncrypted, e.path); err != nil {
			return fmt.Errorf("update %q: %w", e.path, err)
		}
	}

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

	fmt.Printf("Migrated %d entries.\n", len(toMigrate))
	return nil
}

// deriveMEK reads the seal config and derives the MEK from the password.
func deriveMEK(database *sql.DB, password []byte) ([]byte, error) {
	var (
		encryptedMEK    []byte
		salt            []byte
		argTime, argMem uint32
		argThreads      uint8
	)
	err := database.QueryRow(`
		SELECT encrypted_mek, kdf_salt, argon2_time, argon2_memory, argon2_threads
		FROM seal_config WHERE id = 1`).Scan(&encryptedMEK, &salt, &argTime, &argMem, &argThreads)
	if err != nil {
		return nil, fmt.Errorf("read seal config: %w", err)
	}

	params := crypto.Argon2Params{Time: argTime, Memory: argMem, Threads: argThreads}
	kwk := crypto.DeriveKey(password, salt, params)
	defer crypto.Zeroize(kwk)

	mek, err := crypto.Decrypt(kwk, encryptedMEK, nil)
	if err != nil {
		return nil, fmt.Errorf("invalid unseal password")
	}
	return mek, nil
}
