// Package e2ee provides end-to-end encryption operations for the MapleFile SDK.
// This file contains memguard-protected secure memory operations.
package e2ee

import (
	"fmt"

	"github.com/awnumar/memguard"
)

// SecureBuffer wraps memguard.LockedBuffer for type safety
type SecureBuffer struct {
	buffer *memguard.LockedBuffer
}

// NewSecureBuffer creates a new secure buffer from bytes
func NewSecureBuffer(data []byte) (*SecureBuffer, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("cannot create secure buffer from empty data")
	}

	buffer := memguard.NewBufferFromBytes(data)
	return &SecureBuffer{buffer: buffer}, nil
}

// NewSecureBufferRandom creates a new secure buffer with random data
func NewSecureBufferRandom(size int) (*SecureBuffer, error) {
	if size <= 0 {
		return nil, fmt.Errorf("size must be positive")
	}

	buffer := memguard.NewBuffer(size)
	return &SecureBuffer{buffer: buffer}, nil
}

// Bytes returns the underlying bytes (caller must handle carefully)
func (s *SecureBuffer) Bytes() []byte {
	if s.buffer == nil {
		return nil
	}
	return s.buffer.Bytes()
}

// Size returns the size of the buffer
func (s *SecureBuffer) Size() int {
	if s.buffer == nil {
		return 0
	}
	return s.buffer.Size()
}

// Destroy securely destroys the buffer
func (s *SecureBuffer) Destroy() {
	if s.buffer != nil {
		s.buffer.Destroy()
		s.buffer = nil
	}
}

// Copy creates a new SecureBuffer with a copy of the data
func (s *SecureBuffer) Copy() (*SecureBuffer, error) {
	if s.buffer == nil {
		return nil, fmt.Errorf("cannot copy destroyed buffer")
	}

	return NewSecureBuffer(s.buffer.Bytes())
}

// SecureKeyChain is a KeyChain that stores the KEK in protected memory
type SecureKeyChain struct {
	kek          *SecureBuffer // Key Encryption Key in protected memory
	salt         []byte        // Salt (not sensitive, kept in regular memory)
	kdfAlgorithm string        // KDF algorithm used
}

// NewSecureKeyChain creates a new SecureKeyChain with KEK in protected memory.
// This function defaults to Argon2id for backward compatibility.
// For cross-platform compatibility, use NewSecureKeyChainWithAlgorithm instead.
func NewSecureKeyChain(password string, salt []byte) (*SecureKeyChain, error) {
	return NewSecureKeyChainWithAlgorithm(password, salt, Argon2IDAlgorithm)
}

// NewSecureKeyChainWithAlgorithm creates a new SecureKeyChain using the specified KDF algorithm.
// algorithm should be one of: Argon2IDAlgorithm ("argon2id") or PBKDF2Algorithm ("PBKDF2-SHA256").
// The web frontend uses PBKDF2-SHA256, while the native app historically used Argon2id.
func NewSecureKeyChainWithAlgorithm(password string, salt []byte, algorithm string) (*SecureKeyChain, error) {
	// Both algorithms use 16-byte salt
	if len(salt) != 16 {
		return nil, fmt.Errorf("invalid salt size: expected 16, got %d", len(salt))
	}

	// Derive KEK from password using specified algorithm
	kekBytes, err := DeriveKeyFromPasswordWithAlgorithm(password, salt, algorithm)
	if err != nil {
		return nil, fmt.Errorf("failed to derive key from password: %w", err)
	}

	// Store KEK in secure memory immediately
	kek, err := NewSecureBuffer(kekBytes)
	if err != nil {
		ClearBytes(kekBytes)
		return nil, fmt.Errorf("failed to create secure buffer for KEK: %w", err)
	}

	// Clear the temporary KEK bytes
	ClearBytes(kekBytes)

	return &SecureKeyChain{
		kek:          kek,
		salt:         salt,
		kdfAlgorithm: algorithm,
	}, nil
}

// Clear securely clears the SecureKeyChain's sensitive data
func (k *SecureKeyChain) Clear() {
	if k.kek != nil {
		k.kek.Destroy()
		k.kek = nil
	}
}

// DecryptMasterKeySecure decrypts the master key and returns it in a SecureBuffer.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func (k *SecureKeyChain) DecryptMasterKeySecure(encryptedMasterKey *EncryptedKey) (*SecureBuffer, error) {
	if k.kek == nil || k.kek.buffer == nil {
		return nil, fmt.Errorf("keychain has been cleared")
	}

	// Decrypt using KEK from secure memory (auto-detect cipher based on nonce size)
	masterKeyBytes, err := DecryptWithAlgorithm(encryptedMasterKey.Ciphertext, encryptedMasterKey.Nonce, k.kek.Bytes())
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt master key: %w", err)
	}

	// Store decrypted master key in secure memory
	masterKey, err := NewSecureBuffer(masterKeyBytes)
	if err != nil {
		ClearBytes(masterKeyBytes)
		return nil, fmt.Errorf("failed to create secure buffer for master key: %w", err)
	}

	// Clear temporary bytes
	ClearBytes(masterKeyBytes)

	return masterKey, nil
}

// DecryptMasterKey provides backward compatibility by returning []byte.
// For new code, prefer DecryptMasterKeySecure.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func (k *SecureKeyChain) DecryptMasterKey(encryptedMasterKey *EncryptedKey) ([]byte, error) {
	if k.kek == nil || k.kek.buffer == nil {
		return nil, fmt.Errorf("keychain has been cleared")
	}

	// Decrypt using KEK from secure memory (auto-detect cipher)
	return DecryptWithAlgorithm(encryptedMasterKey.Ciphertext, encryptedMasterKey.Nonce, k.kek.Bytes())
}

// EncryptMasterKey encrypts a master key with the KEK using ChaCha20-Poly1305.
// For web frontend compatibility, use EncryptMasterKeySecretBox instead.
func (k *SecureKeyChain) EncryptMasterKey(masterKey []byte) (*EncryptedKey, error) {
	if k.kek == nil || k.kek.buffer == nil {
		return nil, fmt.Errorf("keychain has been cleared")
	}

	encrypted, err := Encrypt(masterKey, k.kek.Bytes())
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt master key: %w", err)
	}

	return &EncryptedKey{
		Ciphertext: encrypted.Ciphertext,
		Nonce:      encrypted.Nonce,
	}, nil
}

// EncryptMasterKeySecretBox encrypts a master key with the KEK using XSalsa20-Poly1305 (SecretBox).
// This is compatible with the web frontend's libsodium implementation.
func (k *SecureKeyChain) EncryptMasterKeySecretBox(masterKey []byte) (*EncryptedKey, error) {
	if k.kek == nil || k.kek.buffer == nil {
		return nil, fmt.Errorf("keychain has been cleared")
	}

	encrypted, err := EncryptWithSecretBox(masterKey, k.kek.Bytes())
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt master key: %w", err)
	}

	return &EncryptedKey{
		Ciphertext: encrypted.Ciphertext,
		Nonce:      encrypted.Nonce,
	}, nil
}

// DecryptPrivateKeySecure decrypts a private key and returns it in a SecureBuffer.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func DecryptPrivateKeySecure(encryptedPrivateKey *EncryptedKey, masterKey *SecureBuffer) (*SecureBuffer, error) {
	if masterKey == nil || masterKey.buffer == nil {
		return nil, fmt.Errorf("master key is nil or destroyed")
	}

	// Decrypt private key (auto-detect cipher based on nonce size)
	privateKeyBytes, err := DecryptWithAlgorithm(encryptedPrivateKey.Ciphertext, encryptedPrivateKey.Nonce, masterKey.Bytes())
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt private key: %w", err)
	}

	// Store in secure memory
	privateKey, err := NewSecureBuffer(privateKeyBytes)
	if err != nil {
		ClearBytes(privateKeyBytes)
		return nil, fmt.Errorf("failed to create secure buffer for private key: %w", err)
	}

	// Clear temporary bytes
	ClearBytes(privateKeyBytes)

	return privateKey, nil
}

// WithSecureBuffer provides a callback pattern for temporary use of secure data
// The buffer is automatically destroyed after the callback returns
func WithSecureBuffer(data []byte, fn func(*SecureBuffer) error) error {
	buf, err := NewSecureBuffer(data)
	if err != nil {
		return err
	}
	defer buf.Destroy()

	return fn(buf)
}

// CopyToSecure copies regular bytes into a new SecureBuffer and clears the source
func CopyToSecure(data []byte) (*SecureBuffer, error) {
	buf, err := NewSecureBuffer(data)
	if err != nil {
		return nil, err
	}

	// Clear the source data
	ClearBytes(data)

	return buf, nil
}