bzzz/pkg/slurp/storage/encrypted_storage.go

package storage

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/anthonyrawlins/bzzz/pkg/crypto"
	"github.com/anthonyrawlins/bzzz/pkg/ucxl"
	slurpContext "github.com/anthonyrawlins/bzzz/pkg/slurp/context"
)

// EncryptedStorageImpl implements the EncryptedStorage interface
type EncryptedStorageImpl struct {
	mu            sync.RWMutex
	crypto        crypto.RoleCrypto
	localStorage  LocalStorage
	keyManager    crypto.KeyManager
	accessControl crypto.AccessController
	auditLogger   crypto.AuditLogger
	metrics       *EncryptionMetrics
}

// EncryptionMetrics tracks encryption-related metrics
type EncryptionMetrics struct {
	mu                    sync.RWMutex
	EncryptOperations     int64
	DecryptOperations     int64
	KeyRotations          int64
	AccessDenials         int64
	EncryptionErrors      int64
	DecryptionErrors      int64
	LastKeyRotation       time.Time
	AverageEncryptTime    time.Duration
	AverageDecryptTime    time.Duration
	ActiveEncryptionKeys  int
	ExpiredKeys           int
}

// NewEncryptedStorage creates a new encrypted storage implementation
func NewEncryptedStorage(
	crypto crypto.RoleCrypto,
	localStorage LocalStorage,
	keyManager crypto.KeyManager,
	accessControl crypto.AccessController,
	auditLogger crypto.AuditLogger,
) *EncryptedStorageImpl {
	return &EncryptedStorageImpl{
		crypto:        crypto,
		localStorage:  localStorage,
		keyManager:    keyManager,
		accessControl: accessControl,
		auditLogger:   auditLogger,
		metrics:       &EncryptionMetrics{},
	}
}

// StoreEncrypted stores data encrypted for specific roles
func (es *EncryptedStorageImpl) StoreEncrypted(
	ctx context.Context,
	key string,
	data interface{},
	roles []string,
) error {
	start := time.Now()
	defer func() {
		es.metrics.mu.Lock()
		es.metrics.EncryptOperations++
		es.updateAverageTime(&es.metrics.AverageEncryptTime, time.Since(start))
		es.metrics.mu.Unlock()
	}()

	// Serialize the data
	dataBytes, err := json.Marshal(data)
	if err != nil {
		es.recordError("encryption", err)
		return fmt.Errorf("failed to marshal data: %w", err)
	}

	// Create encrypted context entries for each role
	for _, role := range roles {
		// Check if role has permission to store data
		if !es.accessControl.CanStore(role, key) {
			es.recordAccessDenial(role, key, "store")
			continue // Skip this role but don't fail the entire operation
		}

		// Create role-specific key
		roleKey := es.generateRoleKey(key, role)

		// Encrypt data for this role
		encryptedData, keyFingerprint, err := es.crypto.EncryptForRole(dataBytes, role)
		if err != nil {
			es.recordError("encryption", err)
			return fmt.Errorf("failed to encrypt data for role %s: %w", role, err)
		}

		// Create encrypted context wrapper
		encCtx := &slurpContext.EncryptedContext{
			Role:           role,
			AccessLevel:    es.determineAccessLevel(role),
			EncryptedData:  encryptedData,
			KeyFingerprint: keyFingerprint,
			CreatedAt:      time.Now(),
		}

		// Store the encrypted context
		storeOptions := &StoreOptions{
			Encrypt:     false, // Already encrypted
			Replicate:   true,
			Index:       true,
			Cache:       true,
			AccessLevel: crypto.AccessLevel(encCtx.AccessLevel),
			Metadata: map[string]interface{}{
				"role":            role,
				"key_fingerprint": keyFingerprint,
				"encrypted":       true,
			},
		}

		if err := es.localStorage.Store(ctx, roleKey, encCtx, storeOptions); err != nil {
			es.recordError("storage", err)
			return fmt.Errorf("failed to store encrypted data for role %s: %w", role, err)
		}

		// Audit log the operation
		es.auditLogger.LogEncryptionOperation(role, key, "store", true)
	}

	return nil
}

// RetrieveDecrypted retrieves and decrypts data for current role
func (es *EncryptedStorageImpl) RetrieveDecrypted(
	ctx context.Context,
	key string,
	role string,
) (interface{}, error) {
	start := time.Now()
	defer func() {
		es.metrics.mu.Lock()
		es.metrics.DecryptOperations++
		es.updateAverageTime(&es.metrics.AverageDecryptTime, time.Since(start))
		es.metrics.mu.Unlock()
	}()

	// Check access permissions
	if !es.accessControl.CanRetrieve(role, key) {
		es.recordAccessDenial(role, key, "retrieve")
		return nil, fmt.Errorf("role %s does not have permission to retrieve key %s", role, key)
	}

	// Generate role-specific key
	roleKey := es.generateRoleKey(key, role)

	// Retrieve encrypted context
	encryptedData, err := es.localStorage.Retrieve(ctx, roleKey)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve encrypted data: %w", err)
	}

	// Cast to encrypted context
	encCtx, ok := encryptedData.(*slurpContext.EncryptedContext)
	if !ok {
		return nil, fmt.Errorf("invalid encrypted context type")
	}

	// Verify role matches
	if encCtx.Role != role {
		es.recordAccessDenial(role, key, "role_mismatch")
		return nil, fmt.Errorf("role mismatch: expected %s, got %s", role, encCtx.Role)
	}

	// Decrypt the data
	decryptedData, err := es.crypto.DecryptForRole(encCtx.EncryptedData, role, encCtx.KeyFingerprint)
	if err != nil {
		es.recordError("decryption", err)
		return nil, fmt.Errorf("failed to decrypt data for role %s: %w", role, err)
	}

	// Deserialize the data
	var result interface{}
	if err := json.Unmarshal(decryptedData, &result); err != nil {
		return nil, fmt.Errorf("failed to unmarshal decrypted data: %w", err)
	}

	// Audit log the operation
	es.auditLogger.LogDecryptionOperation(role, key, "retrieve", true)

	return result, nil
}

// CanAccess checks if a role can access specific data
func (es *EncryptedStorageImpl) CanAccess(
	ctx context.Context,
	key string,
	role string,
) (bool, error) {
	// Check access control rules
	if !es.accessControl.CanRetrieve(role, key) {
		return false, nil
	}

	// Check if encrypted data exists for this role
	roleKey := es.generateRoleKey(key, role)
	exists, err := es.localStorage.Exists(ctx, roleKey)
	if err != nil {
		return false, fmt.Errorf("failed to check existence: %w", err)
	}

	return exists, nil
}

// ListAccessibleKeys lists keys accessible to a role
func (es *EncryptedStorageImpl) ListAccessibleKeys(
	ctx context.Context,
	role string,
) ([]string, error) {
	// List all keys with role prefix
	rolePattern := fmt.Sprintf("role:%s:*", role)
	allKeys, err := es.localStorage.List(ctx, rolePattern)
	if err != nil {
		return nil, fmt.Errorf("failed to list keys for role %s: %w", role, err)
	}

	// Extract original keys from role-prefixed keys
	var accessibleKeys []string
	for _, roleKey := range allKeys {
		originalKey := es.extractOriginalKey(roleKey, role)
		if originalKey != "" && es.accessControl.CanRetrieve(role, originalKey) {
			accessibleKeys = append(accessibleKeys, originalKey)
		}
	}

	return accessibleKeys, nil
}

// ReEncryptForRoles re-encrypts data for different roles
func (es *EncryptedStorageImpl) ReEncryptForRoles(
	ctx context.Context,
	key string,
	newRoles []string,
) error {
	// This requires admin privileges or special re-encryption role
	adminRole := "admin" // TODO: Make this configurable

	// Retrieve original data using admin role
	originalData, err := es.RetrieveDecrypted(ctx, key, adminRole)
	if err != nil {
		return fmt.Errorf("failed to retrieve original data for re-encryption: %w", err)
	}

	// Delete old encrypted versions
	if err := es.deleteAllRoleVersions(ctx, key); err != nil {
		return fmt.Errorf("failed to delete old encrypted versions: %w", err)
	}

	// Store with new roles
	return es.StoreEncrypted(ctx, key, originalData, newRoles)
}

// GetAccessRoles gets roles that can access specific data
func (es *EncryptedStorageImpl) GetAccessRoles(
	ctx context.Context,
	key string,
) ([]string, error) {
	// List all role-prefixed versions of this key
	pattern := fmt.Sprintf("role:*:%s", es.encodeKey(key))
	roleKeys, err := es.localStorage.List(ctx, pattern)
	if err != nil {
		return nil, fmt.Errorf("failed to list role keys: %w", err)
	}

	// Extract roles from keys
	var roles []string
	for _, roleKey := range roleKeys {
		role := es.extractRoleFromKey(roleKey)
		if role != "" {
			roles = append(roles, role)
		}
	}

	return roles, nil
}

// RotateKeys rotates encryption keys
func (es *EncryptedStorageImpl) RotateKeys(
	ctx context.Context,
	maxAge time.Duration,
) error {
	start := time.Now()
	defer func() {
		es.metrics.mu.Lock()
		es.metrics.KeyRotations++
		es.metrics.LastKeyRotation = time.Now()
		es.metrics.mu.Unlock()
	}()

	// Get keys that need rotation
	keysToRotate, err := es.keyManager.GetKeysForRotation(maxAge)
	if err != nil {
		return fmt.Errorf("failed to get keys for rotation: %w", err)
	}

	for _, keyInfo := range keysToRotate {
		// Rotate the key
		if err := es.rotateKey(ctx, keyInfo); err != nil {
			es.recordError("key_rotation", err)
			// Log but continue with other keys
			es.auditLogger.LogKeyRotation(keyInfo.Role, keyInfo.KeyID, false, err.Error())
			continue
		}

		es.auditLogger.LogKeyRotation(keyInfo.Role, keyInfo.KeyID, true, "")
	}

	return nil
}

// ValidateEncryption validates encryption integrity
func (es *EncryptedStorageImpl) ValidateEncryption(
	ctx context.Context,
	key string,
) error {
	// Get all role versions of this key
	roles, err := es.GetAccessRoles(ctx, key)
	if err != nil {
		return fmt.Errorf("failed to get access roles: %w", err)
	}

	// Validate each encrypted version
	for _, role := range roles {
		roleKey := es.generateRoleKey(key, role)
		
		// Retrieve encrypted context
		encryptedData, err := es.localStorage.Retrieve(ctx, roleKey)
		if err != nil {
			return fmt.Errorf("failed to retrieve encrypted data for role %s: %w", role, err)
		}

		encCtx, ok := encryptedData.(*slurpContext.EncryptedContext)
		if !ok {
			return fmt.Errorf("invalid encrypted context type for role %s", role)
		}

		// Validate key fingerprint
		if !es.keyManager.ValidateKeyFingerprint(role, encCtx.KeyFingerprint) {
			return fmt.Errorf("invalid key fingerprint for role %s", role)
		}

		// Try to decrypt (validates encryption integrity)
		_, err = es.crypto.DecryptForRole(encCtx.EncryptedData, role, encCtx.KeyFingerprint)
		if err != nil {
			return fmt.Errorf("encryption validation failed for role %s: %w", role, err)
		}
	}

	return nil
}

// Helper methods

func (es *EncryptedStorageImpl) generateRoleKey(key, role string) string {
	return fmt.Sprintf("role:%s:%s", role, es.encodeKey(key))
}

func (es *EncryptedStorageImpl) encodeKey(key string) string {
	// Use SHA-256 to create a consistent key encoding
	hash := sha256.Sum256([]byte(key))
	return fmt.Sprintf("%x", hash)
}

func (es *EncryptedStorageImpl) extractOriginalKey(roleKey, role string) string {
	prefix := fmt.Sprintf("role:%s:", role)
	if len(roleKey) > len(prefix) && roleKey[:len(prefix)] == prefix {
		return roleKey[len(prefix):]
	}
	return ""
}

func (es *EncryptedStorageImpl) extractRoleFromKey(roleKey string) string {
	// Extract role from "role:ROLE:KEY" format
	if len(roleKey) > 5 && roleKey[:5] == "role:" {
		parts := roleKey[5:] // Remove "role:" prefix
		if idx := len(parts); idx > 0 {
			for i, c := range parts {
				if c == ':' {
					return parts[:i]
				}
			}
		}
	}
	return ""
}

func (es *EncryptedStorageImpl) determineAccessLevel(role string) slurpContext.RoleAccessLevel {
	// Map roles to access levels - this should be configurable
	switch role {
	case "admin", "architect":
		return slurpContext.AccessCritical
	case "senior_developer", "lead":
		return slurpContext.AccessHigh
	case "developer":
		return slurpContext.AccessMedium
	case "junior_developer", "intern":
		return slurpContext.AccessLow
	default:
		return slurpContext.AccessPublic
	}
}

func (es *EncryptedStorageImpl) deleteAllRoleVersions(ctx context.Context, key string) error {
	// Get all roles that have access to this key
	roles, err := es.GetAccessRoles(ctx, key)
	if err != nil {
		return err
	}

	// Delete each role version
	for _, role := range roles {
		roleKey := es.generateRoleKey(key, role)
		if err := es.localStorage.Delete(ctx, roleKey); err != nil {
			// Log but don't fail - may not exist
			es.auditLogger.LogError(fmt.Sprintf("Failed to delete role key %s: %v", roleKey, err))
		}
	}

	return nil
}

func (es *EncryptedStorageImpl) rotateKey(ctx context.Context, keyInfo *crypto.KeyInfo) error {
	// Generate new key for role
	newKeyID, err := es.keyManager.GenerateKey(keyInfo.Role)
	if err != nil {
		return fmt.Errorf("failed to generate new key: %w", err)
	}

	// Re-encrypt all data for this role with new key
	rolePattern := fmt.Sprintf("role:%s:*", keyInfo.Role)
	roleKeys, err := es.localStorage.List(ctx, rolePattern)
	if err != nil {
		return fmt.Errorf("failed to list keys for role %s: %w", keyInfo.Role, err)
	}

	for _, roleKey := range roleKeys {
		if err := es.reEncryptWithNewKey(ctx, roleKey, keyInfo.Role, newKeyID); err != nil {
			// Log but continue
			es.auditLogger.LogError(fmt.Sprintf("Failed to re-encrypt key %s: %v", roleKey, err))
		}
	}

	// Mark old key as deprecated
	return es.keyManager.DeprecateKey(keyInfo.KeyID)
}

func (es *EncryptedStorageImpl) reEncryptWithNewKey(
	ctx context.Context,
	roleKey string,
	role string,
	newKeyID string,
) error {
	// Retrieve and decrypt with old key
	encryptedData, err := es.localStorage.Retrieve(ctx, roleKey)
	if err != nil {
		return err
	}

	encCtx, ok := encryptedData.(*slurpContext.EncryptedContext)
	if !ok {
		return fmt.Errorf("invalid encrypted context type")
	}

	// Decrypt with old key
	decryptedData, err := es.crypto.DecryptForRole(encCtx.EncryptedData, role, encCtx.KeyFingerprint)
	if err != nil {
		return fmt.Errorf("failed to decrypt with old key: %w", err)
	}

	// Encrypt with new key
	newEncryptedData, newKeyFingerprint, err := es.crypto.EncryptForRole(decryptedData, role)
	if err != nil {
		return fmt.Errorf("failed to encrypt with new key: %w", err)
	}

	// Update encrypted context
	encCtx.EncryptedData = newEncryptedData
	encCtx.KeyFingerprint = newKeyFingerprint
	encCtx.CreatedAt = time.Now()

	// Store updated context
	storeOptions := &StoreOptions{
		Encrypt:   false, // Already encrypted
		Replicate: true,
		Index:     true,
		Cache:     true,
		Metadata: map[string]interface{}{
			"role":            role,
			"key_fingerprint": newKeyFingerprint,
			"encrypted":       true,
			"re_encrypted":    true,
		},
	}

	return es.localStorage.Store(ctx, roleKey, encCtx, storeOptions)
}

func (es *EncryptedStorageImpl) recordError(operation string, err error) {
	es.metrics.mu.Lock()
	defer es.metrics.mu.Unlock()

	switch operation {
	case "encryption":
		es.metrics.EncryptionErrors++
	case "decryption":
		es.metrics.DecryptionErrors++
	}

	es.auditLogger.LogError(fmt.Sprintf("%s error: %v", operation, err))
}

func (es *EncryptedStorageImpl) recordAccessDenial(role, key, operation string) {
	es.metrics.mu.Lock()
	es.metrics.AccessDenials++
	es.metrics.mu.Unlock()

	es.auditLogger.LogAccessDenial(role, key, operation)
}

func (es *EncryptedStorageImpl) updateAverageTime(currentAvg *time.Duration, newTime time.Duration) {
	// Simple exponential moving average
	if *currentAvg == 0 {
		*currentAvg = newTime
	} else {
		*currentAvg = time.Duration(float64(*currentAvg)*0.8 + float64(newTime)*0.2)
	}
}

// GetEncryptionMetrics returns encryption metrics
func (es *EncryptedStorageImpl) GetEncryptionMetrics() *EncryptionMetrics {
	es.metrics.mu.RLock()
	defer es.metrics.mu.RUnlock()

	// Return a copy to avoid race conditions
	metricsCopy := *es.metrics
	return &metricsCopy
}