// Package keycache provides secure in-memory caching of cryptographic keys during a session.
// Keys are stored in memguard Enclaves (encrypted at rest in memory) and automatically
// cleared when the application shuts down or the user logs out.
package keycache

import (
	"fmt"
	"sync"

	"github.com/awnumar/memguard"
	"go.uber.org/zap"

	"codeberg.org/mapleopentech/monorepo/native/desktop/maplefile/internal/utils"
)

// Service manages cached cryptographic keys in secure memory
type Service struct {
	logger *zap.Logger
	mu     sync.RWMutex
	// Map of email -> Enclave containing master key
	// Enclave stores data encrypted in memory, must be opened to access
	masterKeys map[string]*memguard.Enclave
}

// ProvideService creates a new key cache service (for Wire)
func ProvideService(logger *zap.Logger) *Service {
	return &Service{
		logger:     logger.Named("keycache"),
		masterKeys: make(map[string]*memguard.Enclave),
	}
}

// StoreMasterKey stores a user's master key in an encrypted memory Enclave
// The key will remain cached until cleared or the app exits
func (s *Service) StoreMasterKey(email string, masterKey []byte) error {
	if email == "" {
		return fmt.Errorf("email is required")
	}

	if len(masterKey) == 0 {
		return fmt.Errorf("master key is empty")
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	// If there's already a cached key, remove it first
	if existing, exists := s.masterKeys[email]; exists {
		// Enclaves are garbage collected when removed from map
		delete(s.masterKeys, email)
		s.logger.Debug("Replaced existing cached master key", zap.String("email", utils.MaskEmail(email)))
		_ = existing // Prevent unused variable warning
	}

	// Create a LockedBuffer from the master key bytes first
	// This locks the memory pages and prevents swapping
	lockedBuf := memguard.NewBufferFromBytes(masterKey)

	// Create an Enclave from the LockedBuffer
	// Enclave stores the data encrypted at rest in memory
	enclave := lockedBuf.Seal()

	// The LockedBuffer is consumed by Seal(), so we don't need to Destroy() it

	// Store the enclave
	s.masterKeys[email] = enclave

	s.logger.Info("Master key cached securely in memory",
		zap.String("email", utils.MaskEmail(email)),
		zap.Int("size", len(masterKey)))

	return nil
}

// GetMasterKey retrieves a cached master key for the given email
// Returns the key bytes and a cleanup function that MUST be called when done
// The cleanup function destroys the LockedBuffer to prevent memory leaks
func (s *Service) GetMasterKey(email string) ([]byte, func(), error) {
	if email == "" {
		return nil, nil, fmt.Errorf("email is required")
	}

	s.mu.RLock()
	enclave, exists := s.masterKeys[email]
	s.mu.RUnlock()

	if !exists {
		return nil, nil, fmt.Errorf("no cached master key found for email: %s", email)
	}

	// Open the enclave to access the master key
	lockedBuf, err := enclave.Open()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to open enclave for reading: %w", err)
	}

	// Get the bytes (caller will use these)
	masterKey := lockedBuf.Bytes()

	// Return cleanup function that destroys the LockedBuffer
	cleanup := func() {
		lockedBuf.Destroy()
	}

	s.logger.Debug("Retrieved cached master key from secure memory",
		zap.String("email", utils.MaskEmail(email)))

	return masterKey, cleanup, nil
}

// WithMasterKey provides a callback pattern for using a cached master key
// The key is automatically cleaned up after the callback returns
func (s *Service) WithMasterKey(email string, fn func([]byte) error) error {
	masterKey, cleanup, err := s.GetMasterKey(email)
	if err != nil {
		return err
	}
	defer cleanup()

	return fn(masterKey)
}

// HasMasterKey checks if a master key is cached for the given email
func (s *Service) HasMasterKey(email string) bool {
	if email == "" {
		return false
	}

	s.mu.RLock()
	defer s.mu.RUnlock()

	_, exists := s.masterKeys[email]
	return exists
}

// ClearMasterKey removes a cached master key for a specific user
func (s *Service) ClearMasterKey(email string) error {
	if email == "" {
		return fmt.Errorf("email is required")
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if enclave, exists := s.masterKeys[email]; exists {
		delete(s.masterKeys, email)
		s.logger.Info("Cleared cached master key from secure memory",
			zap.String("email", utils.MaskEmail(email)))
		_ = enclave // Enclave will be garbage collected
		return nil
	}

	return fmt.Errorf("no cached master key found for email: %s", email)
}

// ClearAll removes all cached master keys
// This should be called on logout or application shutdown
func (s *Service) ClearAll() {
	s.mu.Lock()
	defer s.mu.Unlock()

	count := len(s.masterKeys)
	if count == 0 {
		return
	}

	// Clear all enclaves
	for email := range s.masterKeys {
		delete(s.masterKeys, email)
	}

	s.logger.Info("Cleared all cached master keys from secure memory",
		zap.Int("count", count))
}

// Cleanup performs cleanup operations when the service is shutting down
// This is called by the application shutdown handler
func (s *Service) Cleanup() {
	s.logger.Info("Cleaning up key cache service")
	s.ClearAll()
}