package security

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"math/big"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/accounts/keystore"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/fraktal/mev-beta/internal/logger"
	"golang.org/x/crypto/scrypt"
)

// KeyManager provides secure private key management and transaction signing
type KeyManager struct {
	logger        *logger.Logger
	keystore      *keystore.KeyStore
	encryptionKey []byte
	keys          map[common.Address]*SecureKey
	keysMutex     sync.RWMutex
	config        *KeyManagerConfig
}

// KeyManagerConfig contains configuration for the key manager
type KeyManagerConfig struct {
	KeystorePath    string        // Path to keystore directory
	EncryptionKey   string        // Master encryption key (should come from secure source)
	KeyRotationDays int           // Days before key rotation warning
	MaxSigningRate  int           // Maximum signings per minute
	RequireHardware bool          // Whether to require hardware security module
	BackupPath      string        // Path for encrypted key backups
	AuditLogPath    string        // Path for audit logging
	SessionTimeout  time.Duration // How long before re-authentication required
}

// SecureKey represents a securely stored private key
type SecureKey struct {
	Address         common.Address `json:"address"`
	EncryptedKey    []byte         `json:"encrypted_key"`
	CreatedAt       time.Time      `json:"created_at"`
	LastUsed        time.Time      `json:"last_used"`
	UsageCount      int64          `json:"usage_count"`
	MaxUsage        int64          `json:"max_usage,omitempty"`
	ExpiresAt       *time.Time     `json:"expires_at,omitempty"`
	BackupLocations []string       `json:"backup_locations,omitempty"`
	KeyType         string         `json:"key_type"` // "trading", "emergency", "backup"
	Permissions     KeyPermissions `json:"permissions"`
	IsActive        bool           `json:"is_active"`
}

// KeyPermissions defines what operations a key can perform
type KeyPermissions struct {
	CanSign          bool     `json:"can_sign"`
	CanTransfer      bool     `json:"can_transfer"`
	MaxTransferWei   *big.Int `json:"max_transfer_wei,omitempty"`
	AllowedContracts []string `json:"allowed_contracts,omitempty"`
	RequireConfirm   bool     `json:"require_confirmation"`
}

// SigningRequest represents a request to sign a transaction
type SigningRequest struct {
	Transaction  *types.Transaction
	ChainID      *big.Int
	From         common.Address
	Purpose      string // Description of what this transaction does
	UrgencyLevel int    // 1-5, with 5 being emergency
}

// SigningResult contains the result of a signing operation
type SigningResult struct {
	SignedTx  *types.Transaction
	Signature []byte
	SignedAt  time.Time
	KeyUsed   common.Address
	AuditID   string
	Warnings  []string
}

// AuditEntry represents a security audit log entry
type AuditEntry struct {
	Timestamp  time.Time      `json:"timestamp"`
	Operation  string         `json:"operation"`
	KeyAddress common.Address `json:"key_address"`
	Success    bool           `json:"success"`
	Details    string         `json:"details"`
	IPAddress  string         `json:"ip_address,omitempty"`
	UserAgent  string         `json:"user_agent,omitempty"`
	RiskScore  int            `json:"risk_score"` // 1-10
}

// NewKeyManager creates a new secure key manager
func NewKeyManager(config *KeyManagerConfig, logger *logger.Logger) (*KeyManager, error) {
	if config == nil {
		config = getDefaultConfig()
	}

	// Validate configuration
	if err := validateConfig(config); err != nil {
		return nil, fmt.Errorf("invalid configuration: %w", err)
	}

	// Create keystore directory if it doesn't exist
	if err := os.MkdirAll(config.KeystorePath, 0700); err != nil {
		return nil, fmt.Errorf("failed to create keystore directory: %w", err)
	}

	// Create backup directory if specified
	if config.BackupPath != "" {
		if err := os.MkdirAll(config.BackupPath, 0700); err != nil {
			return nil, fmt.Errorf("failed to create backup directory: %w", err)
		}
	}

	// Initialize keystore
	ks := keystore.NewKeyStore(config.KeystorePath, keystore.StandardScryptN, keystore.StandardScryptP)

	// Derive encryption key from master key
	encryptionKey, err := deriveEncryptionKey(config.EncryptionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to derive encryption key: %w", err)
	}

	km := &KeyManager{
		logger:        logger,
		keystore:      ks,
		encryptionKey: encryptionKey,
		keys:          make(map[common.Address]*SecureKey),
		config:        config,
	}

	// Load existing keys
	if err := km.loadExistingKeys(); err != nil {
		logger.Warn(fmt.Sprintf("Failed to load existing keys: %v", err))
	}

	// Start background tasks
	go km.backgroundTasks()

	logger.Info("Secure key manager initialized")
	return km, nil
}

// GenerateKey creates a new private key with specified permissions
func (km *KeyManager) GenerateKey(keyType string, permissions KeyPermissions) (common.Address, error) {
	// Generate new private key
	privateKey, err := crypto.GenerateKey()
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to generate private key: %w", err)
	}

	address := crypto.PubkeyToAddress(privateKey.PublicKey)

	// Encrypt the private key
	encryptedKey, err := km.encryptPrivateKey(privateKey)
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to encrypt private key: %w", err)
	}

	// Create secure key object
	secureKey := &SecureKey{
		Address:      address,
		EncryptedKey: encryptedKey,
		CreatedAt:    time.Now(),
		LastUsed:     time.Now(),
		UsageCount:   0,
		KeyType:      keyType,
		Permissions:  permissions,
		IsActive:     true, // Mark as active by default
	}

	// Set expiration for certain key types
	if keyType == "emergency" {
		expiresAt := time.Now().Add(30 * 24 * time.Hour) // 30 days
		secureKey.ExpiresAt = &expiresAt
	}

	// Store the key
	km.keysMutex.Lock()
	km.keys[address] = secureKey
	km.keysMutex.Unlock()

	// Create backup
	if err := km.createKeyBackup(secureKey); err != nil {
		km.logger.Warn(fmt.Sprintf("Failed to create backup for key %s: %v", address.Hex(), err))
	}

	// Audit log
	km.auditLog("KEY_GENERATED", address, true, fmt.Sprintf("Generated %s key", keyType))

	km.logger.Info(fmt.Sprintf("Generated new %s key: %s", keyType, address.Hex()))
	return address, nil
}

// ImportKey imports an existing private key
func (km *KeyManager) ImportKey(privateKeyHex string, keyType string, permissions KeyPermissions) (common.Address, error) {
	// Parse private key
	privateKey, err := crypto.HexToECDSA(privateKeyHex)
	if err != nil {
		return common.Address{}, fmt.Errorf("invalid private key: %w", err)
	}

	address := crypto.PubkeyToAddress(privateKey.PublicKey)

	// Check if key already exists
	km.keysMutex.RLock()
	_, exists := km.keys[address]
	km.keysMutex.RUnlock()

	if exists {
		return common.Address{}, fmt.Errorf("key already exists: %s", address.Hex())
	}

	// Encrypt the private key
	encryptedKey, err := km.encryptPrivateKey(privateKey)
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to encrypt private key: %w", err)
	}

	// Create secure key object
	secureKey := &SecureKey{
		Address:      address,
		EncryptedKey: encryptedKey,
		CreatedAt:    time.Now(),
		LastUsed:     time.Now(),
		UsageCount:   0,
		KeyType:      keyType,
		Permissions:  permissions,
		IsActive:     true, // Mark as active by default
	}

	// Store the key
	km.keysMutex.Lock()
	km.keys[address] = secureKey
	km.keysMutex.Unlock()

	// Create backup
	if err := km.createKeyBackup(secureKey); err != nil {
		km.logger.Warn(fmt.Sprintf("Failed to create backup for key %s: %v", address.Hex(), err))
	}

	// Audit log
	km.auditLog("KEY_IMPORTED", address, true, fmt.Sprintf("Imported %s key", keyType))

	km.logger.Info(fmt.Sprintf("Imported %s key: %s", keyType, address.Hex()))
	return address, nil
}

// SignTransaction signs a transaction with comprehensive security checks
func (km *KeyManager) SignTransaction(request *SigningRequest) (*SigningResult, error) {
	// Get the key
	km.keysMutex.RLock()
	secureKey, exists := km.keys[request.From]
	km.keysMutex.RUnlock()

	if !exists {
		km.auditLog("SIGN_FAILED", request.From, false, "Key not found")
		return nil, fmt.Errorf("key not found: %s", request.From.Hex())
	}

	// Security checks
	warnings := make([]string, 0)

	// Check permissions
	if !secureKey.Permissions.CanSign {
		km.auditLog("SIGN_FAILED", request.From, false, "Key not permitted to sign")
		return nil, fmt.Errorf("key %s not permitted to sign transactions", request.From.Hex())
	}

	// Check expiration
	if secureKey.ExpiresAt != nil && time.Now().After(*secureKey.ExpiresAt) {
		km.auditLog("SIGN_FAILED", request.From, false, "Key expired")
		return nil, fmt.Errorf("key %s has expired", request.From.Hex())
	}

	// Check usage limits
	if secureKey.MaxUsage > 0 && secureKey.UsageCount >= secureKey.MaxUsage {
		km.auditLog("SIGN_FAILED", request.From, false, "Usage limit exceeded")
		return nil, fmt.Errorf("key %s usage limit exceeded", request.From.Hex())
	}

	// Check transfer permissions and limits
	if request.Transaction.Value().Sign() > 0 {
		if !secureKey.Permissions.CanTransfer {
			km.auditLog("SIGN_FAILED", request.From, false, "Transfer not permitted")
			return nil, fmt.Errorf("key %s not permitted to transfer value", request.From.Hex())
		}

		if secureKey.Permissions.MaxTransferWei != nil &&
			request.Transaction.Value().Cmp(secureKey.Permissions.MaxTransferWei) > 0 {
			km.auditLog("SIGN_FAILED", request.From, false, "Transfer amount exceeds limit")
			return nil, fmt.Errorf("transfer amount exceeds limit for key %s", request.From.Hex())
		}
	}

	// Check contract interaction permissions
	if request.Transaction.To() != nil {
		contractAddr := request.Transaction.To().Hex()
		if len(secureKey.Permissions.AllowedContracts) > 0 {
			allowed := false
			for _, allowedContract := range secureKey.Permissions.AllowedContracts {
				if contractAddr == allowedContract {
					allowed = true
					break
				}
			}
			if !allowed {
				km.auditLog("SIGN_FAILED", request.From, false, "Contract interaction not permitted")
				return nil, fmt.Errorf("key %s not permitted to interact with contract %s", request.From.Hex(), contractAddr)
			}
		}
	}

	// Rate limiting check
	if err := km.checkRateLimit(request.From); err != nil {
		km.auditLog("SIGN_FAILED", request.From, false, "Rate limit exceeded")
		return nil, fmt.Errorf("rate limit exceeded: %w", err)
	}

	// Warning checks
	if time.Since(secureKey.LastUsed) > 24*time.Hour {
		warnings = append(warnings, "Key has not been used in over 24 hours")
	}

	if secureKey.UsageCount > 1000 {
		warnings = append(warnings, "Key has high usage count - consider rotation")
	}

	// Decrypt private key
	privateKey, err := km.decryptPrivateKey(secureKey.EncryptedKey)
	if err != nil {
		km.auditLog("SIGN_FAILED", request.From, false, "Failed to decrypt private key")
		return nil, fmt.Errorf("failed to decrypt private key: %w", err)
	}
	defer func() {
		// Clear private key from memory
		if privateKey != nil {
			clearPrivateKey(privateKey)
		}
	}()

	// Sign the transaction
	signer := types.NewEIP155Signer(request.ChainID)
	signedTx, err := types.SignTx(request.Transaction, signer, privateKey)
	if err != nil {
		km.auditLog("SIGN_FAILED", request.From, false, "Transaction signing failed")
		return nil, fmt.Errorf("failed to sign transaction: %w", err)
	}

	// Extract signature
	v, r, s := signedTx.RawSignatureValues()
	signature := make([]byte, 65)
	r.FillBytes(signature[0:32])
	s.FillBytes(signature[32:64])
	signature[64] = byte(v.Uint64() - 35 - 2*request.ChainID.Uint64()) // Convert to recovery ID

	// Update key usage
	km.keysMutex.Lock()
	secureKey.LastUsed = time.Now()
	secureKey.UsageCount++
	km.keysMutex.Unlock()

	// Generate audit ID
	auditID := generateAuditID()

	// Create result
	result := &SigningResult{
		SignedTx:  signedTx,
		Signature: signature,
		SignedAt:  time.Now(),
		KeyUsed:   request.From,
		AuditID:   auditID,
		Warnings:  warnings,
	}

	// Audit log
	km.auditLog("TRANSACTION_SIGNED", request.From, true,
		fmt.Sprintf("Signed transaction %s for %s (audit: %s)",
			signedTx.Hash().Hex(), request.Purpose, auditID))

	return result, nil
}

// GetKeyInfo returns information about a key (without sensitive data)
func (km *KeyManager) GetKeyInfo(address common.Address) (*SecureKey, error) {
	km.keysMutex.RLock()
	defer km.keysMutex.RUnlock()

	secureKey, exists := km.keys[address]
	if !exists {
		return nil, fmt.Errorf("key not found: %s", address.Hex())
	}

	// Return a copy without the encrypted key
	info := *secureKey
	info.EncryptedKey = nil

	return &info, nil
}

// ListKeys returns addresses of all managed keys
func (km *KeyManager) ListKeys() []common.Address {
	km.keysMutex.RLock()
	defer km.keysMutex.RUnlock()

	addresses := make([]common.Address, 0, len(km.keys))
	for address := range km.keys {
		addresses = append(addresses, address)
	}

	return addresses
}

// RotateKey creates a new key to replace an existing one
func (km *KeyManager) RotateKey(oldAddress common.Address) (common.Address, error) {
	km.keysMutex.RLock()
	oldKey, exists := km.keys[oldAddress]
	km.keysMutex.RUnlock()

	if !exists {
		return common.Address{}, fmt.Errorf("key not found: %s", oldAddress.Hex())
	}

	// Generate new key with same permissions
	newAddress, err := km.GenerateKey(oldKey.KeyType, oldKey.Permissions)
	if err != nil {
		return common.Address{}, fmt.Errorf("failed to generate new key: %w", err)
	}

	// Mark old key as rotated (don't delete immediately for audit purposes)
	km.keysMutex.Lock()
	oldKey.Permissions.CanSign = false
	oldKey.Permissions.CanTransfer = false
	km.keysMutex.Unlock()

	// Audit log
	km.auditLog("KEY_ROTATED", oldAddress, true,
		fmt.Sprintf("Rotated to new key %s", newAddress.Hex()))

	km.logger.Info(fmt.Sprintf("Rotated key %s to %s", oldAddress.Hex(), newAddress.Hex()))
	return newAddress, nil
}

// encryptPrivateKey encrypts a private key using AES-GCM
func (km *KeyManager) encryptPrivateKey(privateKey *ecdsa.PrivateKey) ([]byte, error) {
	// Convert private key to bytes
	keyBytes := crypto.FromECDSA(privateKey)

	// Create AES cipher
	block, err := aes.NewCipher(km.encryptionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	// Create GCM
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	// Generate nonce
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Encrypt
	ciphertext := gcm.Seal(nonce, nonce, keyBytes, nil)

	// Clear original key bytes
	for i := range keyBytes {
		keyBytes[i] = 0
	}

	return ciphertext, nil
}

// decryptPrivateKey decrypts an encrypted private key
func (km *KeyManager) decryptPrivateKey(encryptedKey []byte) (*ecdsa.PrivateKey, error) {
	// Create AES cipher
	block, err := aes.NewCipher(km.encryptionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	// Create GCM
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	// Extract nonce
	nonceSize := gcm.NonceSize()
	if len(encryptedKey) < nonceSize {
		return nil, fmt.Errorf("encrypted key too short")
	}

	nonce, ciphertext := encryptedKey[:nonceSize], encryptedKey[nonceSize:]

	// Decrypt
	keyBytes, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decryption failed: %w", err)
	}
	defer func() {
		// Clear decrypted bytes
		for i := range keyBytes {
			keyBytes[i] = 0
		}
	}()

	// Convert to private key
	privateKey, err := crypto.ToECDSA(keyBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}

	return privateKey, nil
}

// createKeyBackup creates an encrypted backup of a key
func (km *KeyManager) createKeyBackup(secureKey *SecureKey) error {
	if km.config.BackupPath == "" {
		return nil // Backups not configured
	}

	backupFile := filepath.Join(km.config.BackupPath,
		fmt.Sprintf("key_%s_%d.backup", secureKey.Address.Hex(), time.Now().Unix()))

	// Create backup data
	backupData := struct {
		Address      string    `json:"address"`
		EncryptedKey []byte    `json:"encrypted_key"`
		CreatedAt    time.Time `json:"created_at"`
		KeyType      string    `json:"key_type"`
	}{
		Address:      secureKey.Address.Hex(),
		EncryptedKey: secureKey.EncryptedKey,
		CreatedAt:    secureKey.CreatedAt,
		KeyType:      secureKey.KeyType,
	}

	// Additional encryption for backup
	backupBytes, err := encryptBackupData(backupData, km.encryptionKey)
	if err != nil {
		return fmt.Errorf("failed to encrypt backup: %w", err)
	}

	// Write to file
	if err := os.WriteFile(backupFile, backupBytes, 0600); err != nil {
		return fmt.Errorf("failed to write backup file: %w", err)
	}

	// Update backup locations
	secureKey.BackupLocations = append(secureKey.BackupLocations, backupFile)

	return nil
}

// checkRateLimit checks if signing rate limit is exceeded
func (km *KeyManager) checkRateLimit(address common.Address) error {
	if km.config.MaxSigningRate <= 0 {
		return nil // Rate limiting disabled
	}

	// Implementation would track signing rates per key
	// For now, return nil (rate limiting not implemented)
	return nil
}

// auditLog writes an entry to the audit log
func (km *KeyManager) auditLog(operation string, keyAddress common.Address, success bool, details string) {
	entry := AuditEntry{
		Timestamp:  time.Now(),
		Operation:  operation,
		KeyAddress: keyAddress,
		Success:    success,
		Details:    details,
		RiskScore:  calculateRiskScore(operation, success),
	}

	// Write to audit log
	if km.config.AuditLogPath != "" {
		// Implementation would write to audit log file
		km.logger.Info(fmt.Sprintf("AUDIT: %s %s %v - %s (Risk: %.2f)",
			entry.Operation, entry.KeyAddress.Hex(), entry.Success, entry.Details, entry.RiskScore))
	}
}

// loadExistingKeys loads keys from the keystore
func (km *KeyManager) loadExistingKeys() error {
	// Implementation would load existing keys from storage
	// For now, just log that we're loading
	km.logger.Info("Loading existing keys from keystore")
	return nil
}

// backgroundTasks runs periodic maintenance tasks
func (km *KeyManager) backgroundTasks() {
	ticker := time.NewTicker(1 * time.Hour)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			km.performMaintenance()
		}
	}
}

// performMaintenance performs periodic security maintenance
func (km *KeyManager) performMaintenance() {
	km.keysMutex.RLock()
	defer km.keysMutex.RUnlock()

	now := time.Now()

	for address, key := range km.keys {
		// Check for expired keys
		if key.ExpiresAt != nil && now.After(*key.ExpiresAt) {
			km.logger.Warn(fmt.Sprintf("Key %s has expired", address.Hex()))
		}

		// Check for keys that should be rotated
		if km.config.KeyRotationDays > 0 {
			rotationTime := time.Duration(km.config.KeyRotationDays) * 24 * time.Hour
			if now.Sub(key.CreatedAt) > rotationTime {
				km.logger.Warn(fmt.Sprintf("Key %s should be rotated (age: %v)",
					address.Hex(), now.Sub(key.CreatedAt)))
			}
		}
	}
}

// GetActivePrivateKey returns the active private key for transaction signing
func (km *KeyManager) GetActivePrivateKey() (*ecdsa.PrivateKey, error) {
	// First, check for existing active keys
	km.keysMutex.RLock()
	for address, secureKey := range km.keys {
		if secureKey.IsActive {
			// Decrypt the private key
			privateKey, err := km.decryptPrivateKey(secureKey.EncryptedKey)
			if err != nil {
				km.keysMutex.RUnlock()
				km.auditLog("KEY_DECRYPTION_FAILED", address, false,
					fmt.Sprintf("Failed to decrypt key: %v", err))
				return nil, fmt.Errorf("failed to decrypt active key: %w", err)
			}

			km.keysMutex.RUnlock()
			km.auditLog("KEY_ACCESSED", address, true, "Active private key retrieved")
			return privateKey, nil
		}
	}

	// Check if we need to generate a new key (no keys exist)
	needsNewKey := len(km.keys) == 0
	km.keysMutex.RUnlock()

	// If no active key found and no keys exist, generate a default one
	if needsNewKey {
		km.logger.Info("No keys found, generating default trading key...")
		// Generate a new key pair with default permissions
		defaultPermissions := KeyPermissions{
			CanSign:          true,
			CanTransfer:      true,
			MaxTransferWei:   big.NewInt(1000000000000000000), // 1 ETH max per transaction
			AllowedContracts: []string{},                      // Will be populated with contract addresses
			RequireConfirm:   false,
		}
		km.logger.Info("Calling GenerateKey...")
		address, err := km.GenerateKey("trading", defaultPermissions)
		if err != nil {
			km.logger.Error(fmt.Sprintf("Failed to generate default key: %v", err))
			return nil, fmt.Errorf("failed to generate default key: %w", err)
		}
		km.logger.Info(fmt.Sprintf("Default key generated: %s", address.Hex()))

		// Retrieve the newly generated key
		km.keysMutex.RLock()
		if secureKey, exists := km.keys[address]; exists {
			privateKey, err := km.decryptPrivateKey(secureKey.EncryptedKey)
			km.keysMutex.RUnlock()
			if err != nil {
				return nil, fmt.Errorf("failed to decrypt newly generated key: %w", err)
			}

			km.auditLog("KEY_AUTO_GENERATED", address, true, "Auto-generated active key")
			return privateKey, nil
		}
		km.keysMutex.RUnlock()
	}

	return nil, fmt.Errorf("no active private key available")
}

// Helper functions

func getDefaultConfig() *KeyManagerConfig {
	return &KeyManagerConfig{
		KeystorePath:    "./keystore",
		KeyRotationDays: 90,
		MaxSigningRate:  60, // 60 signings per minute
		RequireHardware: false,
		BackupPath:      "./backups",
		AuditLogPath:    "./audit.log",
		SessionTimeout:  15 * time.Minute,
	}
}

func validateConfig(config *KeyManagerConfig) error {
	if config.KeystorePath == "" {
		return fmt.Errorf("keystore path cannot be empty")
	}
	if config.EncryptionKey == "" {
		return fmt.Errorf("encryption key cannot be empty")
	}
	if len(config.EncryptionKey) < 32 {
		return fmt.Errorf("encryption key must be at least 32 characters")
	}
	return nil
}

func deriveEncryptionKey(masterKey string) ([]byte, error) {
	// Generate secure random salt
	salt := make([]byte, 32)
	if _, err := rand.Read(salt); err != nil {
		return nil, fmt.Errorf("failed to generate random salt: %w", err)
	}

	key, err := scrypt.Key([]byte(masterKey), salt, 32768, 8, 1, 32)
	if err != nil {
		return nil, fmt.Errorf("key derivation failed: %w", err)
	}
	return key, nil
}

func clearPrivateKey(privateKey *ecdsa.PrivateKey) {
	if privateKey != nil && privateKey.D != nil {
		privateKey.D.SetInt64(0)
	}
}

func generateAuditID() string {
	bytes := make([]byte, 16)
	rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

func calculateRiskScore(operation string, success bool) int {
	if !success {
		return 8 // Failed operations are high risk
	}

	switch operation {
	case "TRANSACTION_SIGNED":
		return 3
	case "KEY_GENERATED", "KEY_IMPORTED":
		return 5
	case "KEY_ROTATED":
		return 4
	default:
		return 2
	}
}

func encryptBackupData(data interface{}, key []byte) ([]byte, error) {
	// Convert data to JSON bytes
	jsonData, err := json.Marshal(data)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal backup data: %w", err)
	}

	// Create AES cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
	}

	// Create GCM mode for authenticated encryption
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM mode: %w", err)
	}

	// Generate random nonce
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Encrypt and authenticate the data
	ciphertext := gcm.Seal(nonce, nonce, jsonData, nil)

	return ciphertext, nil
}