package kms_test

import (
	"bytes"
	"testing"
	"time"

	"github.com/jfraeys/fetch_ml/internal/crypto/kms"
)

// TestDEKCache_PutAndGet tests basic cache put and get operations.
func TestDEKCache_PutAndGet(t *testing.T) {
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())
	defer cache.Clear()

	tenantID := "tenant-1"
	artifactID := "artifact-1"
	dek := []byte("test-dek-data-12345678901234567890123456789012")

	// Put DEK in cache
	if err := cache.Put(tenantID, artifactID, "kms-key-1", dek); err != nil {
		t.Fatalf("Put failed: %v", err)
	}

	// Get DEK from cache (KMS available)
	retrieved, ok := cache.Get(tenantID, artifactID, "kms-key-1", false)
	if !ok {
		t.Fatal("Get returned false, expected true")
	}

	if !bytes.Equal(retrieved, dek) {
		t.Error("Retrieved DEK doesn't match original")
	}
}

// TestDEKCache_GetNonexistent tests getting a non-existent entry.
func TestDEKCache_GetNonexistent(t *testing.T) {
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())
	defer cache.Clear()

	_, ok := cache.Get("nonexistent", "nonexistent", "kms-key-1", false)
	if ok {
		t.Error("Get for non-existent key should return false")
	}
}

// TestDEKCache_TTLExpiry tests that entries expire after TTL.
func TestDEKCache_TTLExpiry(t *testing.T) {
	// Use very short TTL for testing
	config := kms.CacheConfig{
		TTL:         50 * time.Millisecond,
		MaxEntries:  100,
		GraceWindow: 100 * time.Millisecond,
	}
	cache := kms.NewDEKCache(config)
	defer cache.Clear()

	tenantID := "tenant-1"
	artifactID := "artifact-1"
	dek := []byte("test-dek-data-12345678901234567890123456789012")

	// Put DEK
	cache.Put(tenantID, artifactID, "kms-key-1", dek)

	// Should be available immediately
	_, ok := cache.Get(tenantID, artifactID, "kms-key-1", false)
	if !ok {
		t.Error("DEK should be available immediately after put")
	}

	// Wait for TTL to expire
	time.Sleep(60 * time.Millisecond)

	// Should not be available after TTL (KMS available)
	_, ok = cache.Get(tenantID, artifactID, "kms-key-1", false)
	if ok {
		t.Error("DEK should not be available after TTL expires")
	}

	// Should be available in grace window (KMS unavailable)
	_, ok = cache.Get(tenantID, artifactID, "kms-key-1", true)
	if !ok {
		t.Error("DEK should be available in grace window when KMS is unavailable")
	}

	// Wait for grace window to expire
	time.Sleep(150 * time.Millisecond)

	// Should not be available after grace window
	_, ok = cache.Get(tenantID, artifactID, "kms-key-1", true)
	if ok {
		t.Error("DEK should not be available after grace window expires")
	}
}

// TestDEKCache_LRUeviction tests LRU eviction when cache is full.
func TestDEKCache_LRUeviction(t *testing.T) {
	config := kms.CacheConfig{
		TTL:         1 * time.Hour, // Long TTL so eviction is due to size
		MaxEntries:  3,
		GraceWindow: 1 * time.Hour,
	}
	cache := kms.NewDEKCache(config)
	defer cache.Clear()

	// Add 3 entries (at capacity)
	for i := 0; i < 3; i++ {
		dek := []byte("dek-data-12345678901234567890123456789012-" + string(rune('0'+i)))
		cache.Put("tenant-1", string(rune('a'+i)), "kms-key-1", dek)
	}

	// Access first entry to make it recently used
	cache.Get("tenant-1", "a", "kms-key-1", false)

	// Add 4th entry (should evict 'b' as it's the oldest unaccessed)
	dek4 := []byte("dek-data-12345678901234567890123456789012-4")
	cache.Put("tenant-1", "d", "kms-key-1", dek4)

	// 'a' should still exist (was accessed)
	_, ok := cache.Get("tenant-1", "a", "kms-key-1", false)
	if !ok {
		t.Error("Entry 'a' should still exist after eviction")
	}

	// 'b' should be evicted
	_, ok = cache.Get("tenant-1", "b", "kms-key-1", false)
	if ok {
		t.Error("Entry 'b' should have been evicted")
	}

	// 'c' and 'd' should exist
	_, ok = cache.Get("tenant-1", "c", "kms-key-1", false)
	if !ok {
		t.Error("Entry 'c' should still exist")
	}
	_, ok = cache.Get("tenant-1", "d", "kms-key-1", false)
	if !ok {
		t.Error("Entry 'd' should exist")
	}
}

// TestDEKCache_Flush tests flushing entries for a specific tenant.
func TestDEKCache_Flush(t *testing.T) {
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())
	defer cache.Clear()

	// Add entries for two tenants
	cache.Put("tenant-1", "artifact-1", "kms-key-1", []byte("dek-1"))
	cache.Put("tenant-1", "artifact-2", "kms-key-1", []byte("dek-2"))
	cache.Put("tenant-2", "artifact-1", "kms-key-2", []byte("dek-3"))

	// Flush tenant-1
	cache.Flush("tenant-1")

	// tenant-1 entries should be gone
	_, ok := cache.Get("tenant-1", "artifact-1", "kms-key-1", false)
	if ok {
		t.Error("tenant-1 artifact-1 should be flushed")
	}
	_, ok = cache.Get("tenant-1", "artifact-2", "kms-key-1", false)
	if ok {
		t.Error("tenant-1 artifact-2 should be flushed")
	}

	// tenant-2 entry should still exist
	_, ok = cache.Get("tenant-2", "artifact-1", "kms-key-2", false)
	if !ok {
		t.Error("tenant-2 artifact-1 should still exist")
	}
}

// TestDEKCache_Clear tests clearing all entries.
func TestDEKCache_Clear(t *testing.T) {
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())

	// Add entries
	cache.Put("tenant-1", "artifact-1", "kms-key-1", []byte("dek-1"))
	cache.Put("tenant-2", "artifact-1", "kms-key-2", []byte("dek-2"))

	// Clear
	cache.Clear()

	// All entries should be gone
	_, ok := cache.Get("tenant-1", "artifact-1", "kms-key-1", false)
	if ok {
		t.Error("All entries should be cleared")
	}
	_, ok = cache.Get("tenant-2", "artifact-1", "kms-key-2", false)
	if ok {
		t.Error("All entries should be cleared")
	}
}

// TestDEKCache_Stats tests cache statistics.
func TestDEKCache_Stats(t *testing.T) {
	config := kms.DefaultCacheConfig()
	cache := kms.NewDEKCache(config)
	defer cache.Clear()

	stats := cache.Stats()

	if stats.Size != 0 {
		t.Errorf("Initial size should be 0, got %d", stats.Size)
	}
	if stats.MaxSize != config.MaxEntries {
		t.Errorf("MaxSize should be %d, got %d", config.MaxEntries, stats.MaxSize)
	}
	if stats.TTL != config.TTL {
		t.Errorf("TTL should be %v, got %v", config.TTL, stats.TTL)
	}
	if stats.GraceWindow != config.GraceWindow {
		t.Errorf("GraceWindow should be %v, got %v", config.GraceWindow, stats.GraceWindow)
	}

	// Add entry
	cache.Put("tenant-1", "artifact-1", "kms-key-1", []byte("dek-1"))

	stats = cache.Stats()
	if stats.Size != 1 {
		t.Errorf("Size should be 1 after put, got %d", stats.Size)
	}
}

// TestDEKCache_EmptyDEK tests that empty DEK is rejected.
func TestDEKCache_EmptyDEK(t *testing.T) {
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())
	defer cache.Clear()

	err := cache.Put("tenant-1", "artifact-1", "kms-key-1", []byte{})
	if err == nil {
		t.Error("Should reject empty DEK")
	}
}

// TestDEKCache_Isolation tests that DEKs are isolated between tenants.
func TestDEKCache_Isolation(t *testing.T) {
	cache := kms.NewDEKCache(kms.DefaultCacheConfig())
	defer cache.Clear()

	// Same artifact ID, different tenants
	cache.Put("tenant-1", "shared-artifact", "kms-key-1", []byte("dek-for-tenant-1"))
	cache.Put("tenant-2", "shared-artifact", "kms-key-2", []byte("dek-for-tenant-2"))

	// Each tenant should get their own DEK
	d1, ok := cache.Get("tenant-1", "shared-artifact", "kms-key-1", false)
	if !ok || string(d1) != "dek-for-tenant-1" {
		t.Error("tenant-1 should get their own DEK")
	}

	d2, ok := cache.Get("tenant-2", "shared-artifact", "kms-key-2", false)
	if !ok || string(d2) != "dek-for-tenant-2" {
		t.Error("tenant-2 should get their own DEK")
	}
}