package crypto

import (
	"sync"
)

// UEKCache holds decrypted User Encryption Keys in memory, keyed by user ID.
// Keys are loaded on login and evicted on logout / token expiry.
//
// This is the trust boundary: UEKs exist in plaintext only in server memory,
// never on disk or in the database.
type UEKCache struct {
	mu    sync.RWMutex
	store map[string][]byte
}

// NewUEKCache creates a new empty cache.
func NewUEKCache() *UEKCache {
	return &UEKCache{store: make(map[string][]byte)}
}

// Store caches a UEK for the given user ID.
func (c *UEKCache) Store(userID string, uek []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	// Copy to avoid caller mutation
	key := make([]byte, len(uek))
	copy(key, uek)
	c.store[userID] = key
}

// Load retrieves the UEK for a user. Returns (uek, true) if present.
func (c *UEKCache) Load(userID string) ([]byte, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	uek, ok := c.store[userID]
	if !ok {
		return nil, false
	}
	// Return a copy
	out := make([]byte, len(uek))
	copy(out, uek)
	return out, true
}

// Evict removes the UEK for a user (logout / session expiry).
func (c *UEKCache) Evict(userID string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	// Zero the memory before deleting
	if uek, ok := c.store[userID]; ok {
		for i := range uek {
			uek[i] = 0
		}
	}
	delete(c.store, userID)
}

// Size returns the number of cached UEKs (for diagnostics).
func (c *UEKCache) Size() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.store)
}