package security

import (
	"encoding/base64"
	"testing"

	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

type encryptedTestRecord struct {
	ID       uint   `gorm:"primaryKey"`
	Secret   string `secure:"access_token"`
	JSONData string `secure:"provider_config"`
}

func TestGORMEncryptionSupportsOnlineKeyRotation(t *testing.T) {
	keyV1 := base64.StdEncoding.EncodeToString([]byte("11111111111111111111111111111111"))
	keyV2 := base64.StdEncoding.EncodeToString([]byte("22222222222222222222222222222222"))
	dsn := "file:gorm-encryption?mode=memory&cache=shared"
	dbV1, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, dbV1.AutoMigrate(&encryptedTestRecord{}))
	encV1, err := NewEncryptor(EncryptionConfig{Enabled: true, AESKey: keyV1, KeyVersion: 1})
	require.NoError(t, err)
	require.NoError(t, RegisterGORMEncryption(dbV1, encV1))

	record := encryptedTestRecord{Secret: "sensitive-value", JSONData: `{"api_key":"provider-secret"}`}
	require.NoError(t, dbV1.Create(&record).Error)
	require.Equal(t, "sensitive-value", record.Secret)
	require.True(t, IsEncrypted(rawEncryptedSecret(t, dbV1, record.ID)))
	require.Contains(t, rawEncryptedSecret(t, dbV1, record.ID), "enc:v1:")
	require.NotContains(t, rawEncryptedJSON(t, dbV1, record.ID), "provider-secret")

	dbV2, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
	require.NoError(t, err)
	encV2, err := NewEncryptorWithPreviousKeys(EncryptionConfig{Enabled: true, AESKey: keyV2, KeyVersion: 2}, map[int]string{1: keyV1})
	require.NoError(t, err)
	require.NoError(t, RegisterGORMEncryption(dbV2, encV2))
	var loaded encryptedTestRecord
	require.NoError(t, dbV2.First(&loaded, record.ID).Error)
	require.Equal(t, "sensitive-value", loaded.Secret)
	require.Equal(t, `{"api_key":"provider-secret"}`, loaded.JSONData)
	require.NoError(t, dbV2.Save(&loaded).Error)
	require.Contains(t, rawEncryptedSecret(t, dbV2, record.ID), "enc:v2:")
}

func rawEncryptedJSON(t *testing.T, db *gorm.DB, id uint) string {
	t.Helper()
	sqlDB, err := db.DB()
	require.NoError(t, err)
	var secret string
	require.NoError(t, sqlDB.QueryRow("SELECT json_data FROM encrypted_test_records WHERE id = ?", id).Scan(&secret))
	return secret
}

func TestEncryptorRejectsPreviousKeyAtCurrentVersion(t *testing.T) {
	key := base64.StdEncoding.EncodeToString([]byte("11111111111111111111111111111111"))
	_, err := NewEncryptorWithPreviousKeys(EncryptionConfig{Enabled: true, AESKey: key, KeyVersion: 2}, map[int]string{2: key})
	require.ErrorContains(t, err, "invalid previous encryption key version")
}

func rawEncryptedSecret(t *testing.T, db *gorm.DB, id uint) string {
	t.Helper()
	sqlDB, err := db.DB()
	require.NoError(t, err)
	var secret string
	require.NoError(t, sqlDB.QueryRow("SELECT secret FROM encrypted_test_records WHERE id = ?", id).Scan(&secret))
	return secret
}