package transit

import (
	"context"
	"encoding/base64"
	"strings"
	"sync"
	"testing"

	"git.wntrmute.dev/mc/metacrypt/internal/barrier"
	"git.wntrmute.dev/mc/metacrypt/internal/engine"
)

// memBarrier is an in-memory barrier for testing.
type memBarrier struct {
	data map[string][]byte
	mu   sync.RWMutex
}

func newMemBarrier() *memBarrier {
	return &memBarrier{data: make(map[string][]byte)}
}

func (m *memBarrier) Unseal(_ []byte) error { return nil }
func (m *memBarrier) Seal() error           { return nil }
func (m *memBarrier) IsSealed() bool        { return false }

func (m *memBarrier) Get(_ context.Context, path string) ([]byte, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	v, ok := m.data[path]
	if !ok {
		return nil, barrier.ErrNotFound
	}
	cp := make([]byte, len(v))
	copy(cp, v)
	return cp, nil
}

func (m *memBarrier) Put(_ context.Context, path string, value []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	cp := make([]byte, len(value))
	copy(cp, value)
	m.data[path] = cp
	return nil
}

func (m *memBarrier) Delete(_ context.Context, path string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.data, path)
	return nil
}

func (m *memBarrier) List(_ context.Context, prefix string) ([]string, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var paths []string
	for k := range m.data {
		if strings.HasPrefix(k, prefix) {
			paths = append(paths, strings.TrimPrefix(k, prefix))
		}
	}
	return paths, nil
}

func adminCaller() *engine.CallerInfo {
	return &engine.CallerInfo{Username: "admin", Roles: []string{"admin"}, IsAdmin: true}
}

func userCaller() *engine.CallerInfo {
	return &engine.CallerInfo{Username: "user", Roles: []string{"user"}, IsAdmin: false}
}

func guestCaller() *engine.CallerInfo {
	return &engine.CallerInfo{Username: "guest", Roles: []string{"guest"}, IsAdmin: false}
}

func setupEngine(t *testing.T) (*TransitEngine, *memBarrier) {
	t.Helper()
	b := newMemBarrier()
	eng := NewTransitEngine()
	ctx := context.Background()
	mountPath := "engine/transit/test/"

	err := eng.Initialize(ctx, b, mountPath, nil)
	if err != nil {
		t.Fatalf("Initialize: %v", err)
	}

	te := eng.(*TransitEngine)
	return te, b
}

func setupEngineWithUnseal(t *testing.T) (*TransitEngine, *memBarrier) {
	t.Helper()
	te, b := setupEngine(t)

	// Seal and unseal to test the full lifecycle.
	if err := te.Seal(); err != nil {
		t.Fatalf("Seal: %v", err)
	}

	eng2 := NewTransitEngine()
	ctx := context.Background()
	if err := eng2.Unseal(ctx, b, "engine/transit/test/"); err != nil {
		t.Fatalf("Unseal: %v", err)
	}
	return eng2.(*TransitEngine), b
}

func createKey(t *testing.T, te *TransitEngine, name, keyType string) {
	t.Helper()
	ctx := context.Background()
	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "create-key",
		CallerInfo: adminCaller(),
		Data: map[string]interface{}{
			"name": name,
			"type": keyType,
		},
	})
	if err != nil {
		t.Fatalf("create-key %s: %v", name, err)
	}
}

func TestInitializeAndUnseal(t *testing.T) {
	te, b := setupEngine(t)

	// Create a key so unseal has something to load.
	createKey(t, te, "mykey", "aes256-gcm")

	if err := te.Seal(); err != nil {
		t.Fatalf("Seal: %v", err)
	}

	eng2 := NewTransitEngine()
	if err := eng2.Unseal(context.Background(), b, "engine/transit/test/"); err != nil {
		t.Fatalf("Unseal: %v", err)
	}

	te2 := eng2.(*TransitEngine)
	if _, ok := te2.keys["mykey"]; !ok {
		t.Fatal("expected key 'mykey' after unseal")
	}
}

func TestCreateKeyAllTypes(t *testing.T) {
	types := []string{"aes256-gcm", "chacha20-poly", "ed25519", "ecdsa-p256", "ecdsa-p384", "hmac-sha256", "hmac-sha512"}
	for _, kt := range types {
		t.Run(kt, func(t *testing.T) {
			te, _ := setupEngine(t)
			createKey(t, te, "key-"+kt, kt)

			ks, ok := te.keys["key-"+kt]
			if !ok {
				t.Fatal("key not created")
			}
			if ks.config.CurrentVersion != 1 {
				t.Fatalf("expected version 1, got %d", ks.config.CurrentVersion)
			}
			if ks.config.MinDecryptionVersion != 1 {
				t.Fatalf("expected min_decryption_version 1, got %d", ks.config.MinDecryptionVersion)
			}
		})
	}
}

func TestEncryptDecryptAES(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "aes-key", "aes256-gcm")
	ctx := context.Background()

	plaintext := base64.StdEncoding.EncodeToString([]byte("hello world"))

	// Encrypt.
	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data: map[string]interface{}{
			"key":       "aes-key",
			"plaintext": plaintext,
		},
	})
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}

	ct, _ := resp.Data["ciphertext"].(string)
	if !strings.HasPrefix(ct, "metacrypt:v1:") {
		t.Fatalf("unexpected ciphertext format: %s", ct)
	}

	// Decrypt.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "decrypt",
		CallerInfo: userCaller(),
		Data: map[string]interface{}{
			"key":        "aes-key",
			"ciphertext": ct,
		},
	})
	if err != nil {
		t.Fatalf("decrypt: %v", err)
	}

	ptB64, _ := resp.Data["plaintext"].(string)
	decoded, _ := base64.StdEncoding.DecodeString(ptB64)
	if string(decoded) != "hello world" {
		t.Fatalf("expected 'hello world', got %q", string(decoded))
	}
}

func TestEncryptDecryptChaCha(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "chacha-key", "chacha20-poly")
	ctx := context.Background()

	plaintext := base64.StdEncoding.EncodeToString([]byte("secret data"))

	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "chacha-key", "plaintext": plaintext},
	})
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}

	ct, _ := resp.Data["ciphertext"].(string)
	if !strings.HasPrefix(ct, "metacrypt:v1:") {
		t.Fatalf("unexpected ciphertext format: %s", ct)
	}

	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "chacha-key", "ciphertext": ct},
	})
	if err != nil {
		t.Fatalf("decrypt: %v", err)
	}

	ptB64, _ := resp.Data["plaintext"].(string)
	decoded, _ := base64.StdEncoding.DecodeString(ptB64)
	if string(decoded) != "secret data" {
		t.Fatalf("expected 'secret data', got %q", string(decoded))
	}
}

func TestEncryptWithContext(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "ctx-key", "aes256-gcm")
	ctx := context.Background()

	plaintext := base64.StdEncoding.EncodeToString([]byte("context test"))
	aadB64 := base64.StdEncoding.EncodeToString([]byte("my-context"))

	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "ctx-key", "plaintext": plaintext, "context": aadB64},
	})
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}
	ct, _ := resp.Data["ciphertext"].(string)

	// Decrypt with correct context.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "ctx-key", "ciphertext": ct, "context": aadB64},
	})
	if err != nil {
		t.Fatalf("decrypt: %v", err)
	}

	// Decrypt with wrong context should fail.
	wrongCtx := base64.StdEncoding.EncodeToString([]byte("wrong-context"))
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "ctx-key", "ciphertext": ct, "context": wrongCtx},
	})
	if err == nil {
		t.Fatal("expected decrypt to fail with wrong context")
	}
}

func TestKeyRotationAndDecrypt(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "rot-key", "aes256-gcm")
	ctx := context.Background()

	plaintext := base64.StdEncoding.EncodeToString([]byte("rotate me"))

	// Encrypt with v1.
	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "rot-key", "plaintext": plaintext},
	})
	if err != nil {
		t.Fatalf("encrypt v1: %v", err)
	}
	ctV1, _ := resp.Data["ciphertext"].(string)

	// Rotate.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "rotate-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "rot-key"},
	})
	if err != nil {
		t.Fatalf("rotate: %v", err)
	}

	// Encrypt with v2.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "rot-key", "plaintext": plaintext},
	})
	if err != nil {
		t.Fatalf("encrypt v2: %v", err)
	}
	ctV2, _ := resp.Data["ciphertext"].(string)

	if !strings.HasPrefix(ctV2, "metacrypt:v2:") {
		t.Fatalf("expected v2 ciphertext, got %s", ctV2)
	}

	// Both ciphertexts should decrypt.
	for _, ct := range []string{ctV1, ctV2} {
		resp, err = te.HandleRequest(ctx, &engine.Request{
			Operation:  "decrypt",
			CallerInfo: userCaller(),
			Data:       map[string]interface{}{"key": "rot-key", "ciphertext": ct},
		})
		if err != nil {
			t.Fatalf("decrypt: %v", err)
		}
		ptB64, _ := resp.Data["plaintext"].(string)
		decoded, _ := base64.StdEncoding.DecodeString(ptB64)
		if string(decoded) != "rotate me" {
			t.Fatalf("expected 'rotate me', got %q", string(decoded))
		}
	}
}

func TestUpdateKeyConfig(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "cfg-key", "aes256-gcm")
	ctx := context.Background()

	plaintext := base64.StdEncoding.EncodeToString([]byte("old data"))

	// Encrypt with v1.
	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "cfg-key", "plaintext": plaintext},
	})
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}
	ctV1, _ := resp.Data["ciphertext"].(string)

	// Rotate to v2.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "rotate-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "cfg-key"},
	})
	if err != nil {
		t.Fatalf("rotate: %v", err)
	}

	// Advance min_decryption_version to 2.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "update-key-config",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "cfg-key", "min_decryption_version": float64(2)},
	})
	if err != nil {
		t.Fatalf("update-key-config: %v", err)
	}

	// v1 ciphertext should be rejected.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "cfg-key", "ciphertext": ctV1},
	})
	if err == nil {
		t.Fatal("expected decrypt to fail for v1 ciphertext")
	}

	// Cannot decrease min_decryption_version.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "update-key-config",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "cfg-key", "min_decryption_version": float64(1)},
	})
	if err == nil {
		t.Fatal("expected error decreasing min_decryption_version")
	}

	// Cannot exceed current version.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "update-key-config",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "cfg-key", "min_decryption_version": float64(99)},
	})
	if err == nil {
		t.Fatal("expected error exceeding current version")
	}
}

func TestTrimKey(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "trim-key", "aes256-gcm")
	ctx := context.Background()

	// Rotate to v2, v3.
	for i := 0; i < 2; i++ {
		_, err := te.HandleRequest(ctx, &engine.Request{
			Operation:  "rotate-key",
			CallerInfo: adminCaller(),
			Data:       map[string]interface{}{"name": "trim-key"},
		})
		if err != nil {
			t.Fatalf("rotate: %v", err)
		}
	}

	// Set min_decryption_version to 2.
	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "update-key-config",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "trim-key", "min_decryption_version": float64(2)},
	})
	if err != nil {
		t.Fatalf("update-key-config: %v", err)
	}

	// Trim.
	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "trim-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "trim-key"},
	})
	if err != nil {
		t.Fatalf("trim-key: %v", err)
	}

	trimmed, _ := resp.Data["trimmed"].(int)
	if trimmed != 1 {
		t.Fatalf("expected 1 trimmed, got %d", trimmed)
	}

	// Version 1 should be gone.
	if _, ok := te.keys["trim-key"].versions[1]; ok {
		t.Fatal("version 1 should have been trimmed")
	}
	if _, ok := te.keys["trim-key"].versions[2]; !ok {
		t.Fatal("version 2 should still exist")
	}
}

func TestSignVerifyEd25519(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "ed-key", "ed25519")
	ctx := context.Background()

	input := base64.StdEncoding.EncodeToString([]byte("sign this"))

	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "sign",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "ed-key", "input": input},
	})
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	sig, _ := resp.Data["signature"].(string)
	if !strings.HasPrefix(sig, "metacrypt:v1:") {
		t.Fatalf("unexpected signature format: %s", sig)
	}

	// Verify.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "verify",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "ed-key", "input": input, "signature": sig},
	})
	if err != nil {
		t.Fatalf("verify: %v", err)
	}
	valid, _ := resp.Data["valid"].(bool)
	if !valid {
		t.Fatal("expected valid signature")
	}

	// Wrong input should fail verification.
	wrongInput := base64.StdEncoding.EncodeToString([]byte("wrong data"))
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "verify",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "ed-key", "input": wrongInput, "signature": sig},
	})
	if err != nil {
		t.Fatalf("verify: %v", err)
	}
	valid, _ = resp.Data["valid"].(bool)
	if valid {
		t.Fatal("expected invalid signature for wrong input")
	}
}

func TestSignVerifyECDSA(t *testing.T) {
	for _, keyType := range []string{"ecdsa-p256", "ecdsa-p384"} {
		t.Run(keyType, func(t *testing.T) {
			te, _ := setupEngine(t)
			createKey(t, te, "ec-key", keyType)
			ctx := context.Background()

			input := base64.StdEncoding.EncodeToString([]byte("ecdsa test"))

			resp, err := te.HandleRequest(ctx, &engine.Request{
				Operation:  "sign",
				CallerInfo: userCaller(),
				Data:       map[string]interface{}{"key": "ec-key", "input": input},
			})
			if err != nil {
				t.Fatalf("sign: %v", err)
			}
			sig, _ := resp.Data["signature"].(string)

			resp, err = te.HandleRequest(ctx, &engine.Request{
				Operation:  "verify",
				CallerInfo: userCaller(),
				Data:       map[string]interface{}{"key": "ec-key", "input": input, "signature": sig},
			})
			if err != nil {
				t.Fatalf("verify: %v", err)
			}
			valid, _ := resp.Data["valid"].(bool)
			if !valid {
				t.Fatal("expected valid signature")
			}
		})
	}
}

func TestSignRejectsSymmetricAndHMAC(t *testing.T) {
	for _, keyType := range []string{"aes256-gcm", "hmac-sha256"} {
		t.Run(keyType, func(t *testing.T) {
			te, _ := setupEngine(t)
			createKey(t, te, "sym-key", keyType)
			ctx := context.Background()

			input := base64.StdEncoding.EncodeToString([]byte("test"))
			_, err := te.HandleRequest(ctx, &engine.Request{
				Operation:  "sign",
				CallerInfo: userCaller(),
				Data:       map[string]interface{}{"key": "sym-key", "input": input},
			})
			if err == nil {
				t.Fatal("expected error signing with non-asymmetric key")
			}
		})
	}
}

func TestHMACComputeAndVerify(t *testing.T) {
	for _, keyType := range []string{"hmac-sha256", "hmac-sha512"} {
		t.Run(keyType, func(t *testing.T) {
			te, _ := setupEngine(t)
			createKey(t, te, "hmac-key", keyType)
			ctx := context.Background()

			input := base64.StdEncoding.EncodeToString([]byte("hmac me"))

			// Compute.
			resp, err := te.HandleRequest(ctx, &engine.Request{
				Operation:  "hmac",
				CallerInfo: userCaller(),
				Data:       map[string]interface{}{"key": "hmac-key", "input": input},
			})
			if err != nil {
				t.Fatalf("hmac compute: %v", err)
			}
			hmacStr, _ := resp.Data["hmac"].(string)
			if !strings.HasPrefix(hmacStr, "metacrypt:v1:") {
				t.Fatalf("unexpected hmac format: %s", hmacStr)
			}

			// Verify.
			resp, err = te.HandleRequest(ctx, &engine.Request{
				Operation:  "hmac",
				CallerInfo: userCaller(),
				Data:       map[string]interface{}{"key": "hmac-key", "input": input, "hmac": hmacStr},
			})
			if err != nil {
				t.Fatalf("hmac verify: %v", err)
			}
			valid, _ := resp.Data["valid"].(bool)
			if !valid {
				t.Fatal("expected valid HMAC")
			}

			// Wrong input should fail.
			wrongInput := base64.StdEncoding.EncodeToString([]byte("wrong data"))
			resp, err = te.HandleRequest(ctx, &engine.Request{
				Operation:  "hmac",
				CallerInfo: userCaller(),
				Data:       map[string]interface{}{"key": "hmac-key", "input": wrongInput, "hmac": hmacStr},
			})
			if err != nil {
				t.Fatalf("hmac verify wrong: %v", err)
			}
			valid, _ = resp.Data["valid"].(bool)
			if valid {
				t.Fatal("expected invalid HMAC for wrong input")
			}
		})
	}
}

func TestBatchEncryptDecrypt(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "batch-key", "aes256-gcm")
	ctx := context.Background()

	items := []interface{}{
		map[string]interface{}{
			"plaintext": base64.StdEncoding.EncodeToString([]byte("item1")),
			"reference": "ref1",
		},
		map[string]interface{}{
			"plaintext": base64.StdEncoding.EncodeToString([]byte("item2")),
			"reference": "ref2",
		},
	}

	// Batch encrypt.
	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "batch-encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "batch-key", "items": items},
	})
	if err != nil {
		t.Fatalf("batch-encrypt: %v", err)
	}

	results, _ := resp.Data["results"].([]interface{})
	if len(results) != 2 {
		t.Fatalf("expected 2 results, got %d", len(results))
	}

	// Build decrypt items.
	decryptItems := make([]interface{}, len(results))
	for i, r := range results {
		br, ok := r.(batchResult)
		if !ok {
			t.Fatalf("expected batchResult, got %T", r)
		}
		if br.Error != "" {
			t.Fatalf("batch encrypt item %d error: %s", i, br.Error)
		}
		decryptItems[i] = map[string]interface{}{
			"ciphertext": br.Ciphertext,
			"reference":  br.Reference,
		}
	}

	// Batch decrypt.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "batch-decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "batch-key", "items": decryptItems},
	})
	if err != nil {
		t.Fatalf("batch-decrypt: %v", err)
	}

	results, _ = resp.Data["results"].([]interface{})
	expected := []string{"item1", "item2"}
	for i, r := range results {
		br, ok := r.(batchResult)
		if !ok {
			t.Fatalf("expected batchResult, got %T", r)
		}
		if br.Error != "" {
			t.Fatalf("batch decrypt item %d error: %s", i, br.Error)
		}
		decoded, _ := base64.StdEncoding.DecodeString(br.Plaintext)
		if string(decoded) != expected[i] {
			t.Fatalf("item %d: expected %q, got %q", i, expected[i], string(decoded))
		}
	}
}

func TestBatchPartialErrors(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "batch-err-key", "aes256-gcm")
	ctx := context.Background()

	items := []interface{}{
		map[string]interface{}{
			"ciphertext": "metacrypt:v1:invalidbase64!!!",
			"reference":  "bad",
		},
		map[string]interface{}{
			"ciphertext": "metacrypt:v1:" + base64.StdEncoding.EncodeToString([]byte("not-valid-ciphertext")),
			"reference":  "also-bad",
		},
	}

	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "batch-decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "batch-err-key", "items": items},
	})
	if err != nil {
		t.Fatalf("batch-decrypt: %v", err)
	}

	results, _ := resp.Data["results"].([]interface{})
	if len(results) != 2 {
		t.Fatalf("expected 2 results, got %d", len(results))
	}
	for i, r := range results {
		br, ok := r.(batchResult)
		if !ok {
			t.Fatalf("expected batchResult, got %T", r)
		}
		if br.Error == "" {
			t.Fatalf("item %d: expected error", i)
		}
	}
}

func TestRewrap(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "rewrap-key", "aes256-gcm")
	ctx := context.Background()

	plaintext := base64.StdEncoding.EncodeToString([]byte("rewrap me"))

	// Encrypt with v1.
	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "rewrap-key", "plaintext": plaintext},
	})
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}
	ctV1, _ := resp.Data["ciphertext"].(string)

	// Rotate.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "rotate-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "rewrap-key"},
	})
	if err != nil {
		t.Fatalf("rotate: %v", err)
	}

	// Rewrap.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "rewrap",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "rewrap-key", "ciphertext": ctV1},
	})
	if err != nil {
		t.Fatalf("rewrap: %v", err)
	}
	ctV2, _ := resp.Data["ciphertext"].(string)
	if !strings.HasPrefix(ctV2, "metacrypt:v2:") {
		t.Fatalf("expected v2 ciphertext after rewrap, got %s", ctV2)
	}

	// Decrypt rewrapped ciphertext.
	resp, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "decrypt",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "rewrap-key", "ciphertext": ctV2},
	})
	if err != nil {
		t.Fatalf("decrypt rewrapped: %v", err)
	}
	ptB64, _ := resp.Data["plaintext"].(string)
	decoded, _ := base64.StdEncoding.DecodeString(ptB64)
	if string(decoded) != "rewrap me" {
		t.Fatalf("expected 'rewrap me', got %q", string(decoded))
	}
}

func TestAuthEnforcement(t *testing.T) {
	te, _ := setupEngine(t)
	ctx := context.Background()

	// Admin-only operations should fail for users.
	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "create-key",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"name": "test", "type": "aes256-gcm"},
	})
	if err == nil {
		t.Fatal("expected create-key to fail for user")
	}

	// Admin-only operations should fail for guests.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "create-key",
		CallerInfo: guestCaller(),
		Data:       map[string]interface{}{"name": "test", "type": "aes256-gcm"},
	})
	if err == nil {
		t.Fatal("expected create-key to fail for guest")
	}

	// User operations should fail for guests.
	createKey(t, te, "auth-key", "aes256-gcm")

	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "list-keys",
		CallerInfo: guestCaller(),
	})
	if err == nil {
		t.Fatal("expected list-keys to fail for guest")
	}

	// Encrypt should fail for guest.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "encrypt",
		CallerInfo: guestCaller(),
		Data:       map[string]interface{}{"key": "auth-key", "plaintext": base64.StdEncoding.EncodeToString([]byte("test"))},
	})
	if err == nil {
		t.Fatal("expected encrypt to fail for guest")
	}

	// Unauthenticated should fail.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "create-key",
		CallerInfo: nil,
		Data:       map[string]interface{}{"name": "test", "type": "aes256-gcm"},
	})
	if err == nil {
		t.Fatal("expected create-key to fail without auth")
	}
}

func TestDeleteKeyWithAndWithoutAllowDeletion(t *testing.T) {
	te, _ := setupEngine(t)
	ctx := context.Background()

	createKey(t, te, "nodelete-key", "aes256-gcm")

	// Should fail: allow_deletion is false.
	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "delete-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "nodelete-key"},
	})
	if err == nil {
		t.Fatal("expected delete to fail without allow_deletion")
	}

	// Enable deletion.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "update-key-config",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "nodelete-key", "allow_deletion": true},
	})
	if err != nil {
		t.Fatalf("update-key-config: %v", err)
	}

	// Should succeed now.
	_, err = te.HandleRequest(ctx, &engine.Request{
		Operation:  "delete-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "nodelete-key"},
	})
	if err != nil {
		t.Fatalf("delete-key: %v", err)
	}

	if _, ok := te.keys["nodelete-key"]; ok {
		t.Fatal("key should have been deleted")
	}
}

func TestGetPublicKey(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "pubkey", "ed25519")
	ctx := context.Background()

	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "get-public-key",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"name": "pubkey"},
	})
	if err != nil {
		t.Fatalf("get-public-key: %v", err)
	}
	pk, _ := resp.Data["public_key"].(string)
	if pk == "" {
		t.Fatal("expected non-empty public key")
	}
	ver, _ := resp.Data["version"].(int)
	if ver != 1 {
		t.Fatalf("expected version 1, got %d", ver)
	}
}

func TestGetPublicKeyRejectsSymmetric(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "sym", "aes256-gcm")
	ctx := context.Background()

	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "get-public-key",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"name": "sym"},
	})
	if err == nil {
		t.Fatal("expected error getting public key for symmetric key")
	}
}

func TestBatchRewrap(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "brwrap-key", "aes256-gcm")
	ctx := context.Background()

	pt1 := base64.StdEncoding.EncodeToString([]byte("item1"))
	pt2 := base64.StdEncoding.EncodeToString([]byte("item2"))

	// Encrypt two items.
	var cts []string
	for _, pt := range []string{pt1, pt2} {
		resp, _ := te.HandleRequest(ctx, &engine.Request{
			Operation:  "encrypt",
			CallerInfo: userCaller(),
			Data:       map[string]interface{}{"key": "brwrap-key", "plaintext": pt},
		})
		ct, _ := resp.Data["ciphertext"].(string)
		cts = append(cts, ct)
	}

	// Rotate.
	te.HandleRequest(ctx, &engine.Request{
		Operation:  "rotate-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "brwrap-key"},
	})

	// Batch rewrap.
	items := []interface{}{
		map[string]interface{}{"ciphertext": cts[0], "reference": "r1"},
		map[string]interface{}{"ciphertext": cts[1], "reference": "r2"},
	}

	resp, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "batch-rewrap",
		CallerInfo: userCaller(),
		Data:       map[string]interface{}{"key": "brwrap-key", "items": items},
	})
	if err != nil {
		t.Fatalf("batch-rewrap: %v", err)
	}

	results, _ := resp.Data["results"].([]interface{})
	for i, r := range results {
		br, ok := r.(batchResult)
		if !ok {
			t.Fatalf("expected batchResult, got %T", r)
		}
		if br.Error != "" {
			t.Fatalf("item %d error: %s", i, br.Error)
		}
		if !strings.HasPrefix(br.Ciphertext, "metacrypt:v2:") {
			t.Fatalf("item %d: expected v2 ciphertext, got %s", i, br.Ciphertext)
		}
	}
}

func TestDuplicateKeyCreation(t *testing.T) {
	te, _ := setupEngine(t)
	createKey(t, te, "dup-key", "aes256-gcm")
	ctx := context.Background()

	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "create-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "dup-key", "type": "aes256-gcm"},
	})
	if err == nil {
		t.Fatal("expected error creating duplicate key")
	}
}

func TestInvalidKeyType(t *testing.T) {
	te, _ := setupEngine(t)
	ctx := context.Background()

	_, err := te.HandleRequest(ctx, &engine.Request{
		Operation:  "create-key",
		CallerInfo: adminCaller(),
		Data:       map[string]interface{}{"name": "bad", "type": "invalid-type"},
	})
	if err == nil {
		t.Fatal("expected error for invalid key type")
	}
}