// Package e2ee provides end-to-end encryption operations for the MapleFile SDK.
package e2ee

import (
	"fmt"
)

// KeyChain holds the key encryption key derived from the user's password.
// It provides methods for decrypting keys in the E2EE chain.
type KeyChain struct {
	kek          []byte // Key Encryption Key derived from password
	salt         []byte // Password salt used for key derivation
	kdfAlgorithm string // KDF algorithm used ("argon2id" or "PBKDF2-SHA256")
}

// EncryptedKey represents a key encrypted with another key.
type EncryptedKey struct {
	Ciphertext []byte `json:"ciphertext"`
	Nonce      []byte `json:"nonce"`
}

// NewKeyChain creates a new KeyChain by deriving the KEK from the password and salt.
// This function defaults to Argon2id for backward compatibility.
// For cross-platform compatibility, use NewKeyChainWithAlgorithm instead.
func NewKeyChain(password string, salt []byte) (*KeyChain, error) {
	return NewKeyChainWithAlgorithm(password, salt, Argon2IDAlgorithm)
}

// NewKeyChainWithAlgorithm creates a new KeyChain 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 NewKeyChainWithAlgorithm(password string, salt []byte, algorithm string) (*KeyChain, error) {
	// Validate salt size (both algorithms use 16-byte salt)
	if len(salt) != 16 {
		return nil, fmt.Errorf("invalid salt size: expected 16, got %d", len(salt))
	}

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

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

// Clear securely clears the KeyChain's sensitive data from memory.
// This should be called when the KeyChain is no longer needed.
func (k *KeyChain) Clear() {
	if k.kek != nil {
		ClearBytes(k.kek)
		k.kek = nil
	}
}

// DecryptMasterKey decrypts the user's master key using the KEK.
// This method auto-detects the cipher based on nonce size:
// - 12-byte nonce: ChaCha20-Poly1305 (native app)
// - 24-byte nonce: XSalsa20-Poly1305 (web frontend)
func (k *KeyChain) DecryptMasterKey(encryptedMasterKey *EncryptedKey) ([]byte, error) {
	if k.kek == nil {
		return nil, fmt.Errorf("keychain has been cleared")
	}

	// Auto-detect cipher based on nonce size
	masterKey, err := DecryptWithAlgorithm(encryptedMasterKey.Ciphertext, encryptedMasterKey.Nonce, k.kek)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt master key: %w", err)
	}

	return masterKey, nil
}

// DecryptCollectionKey decrypts a collection key using the master key.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func DecryptCollectionKey(encryptedCollectionKey *EncryptedKey, masterKey []byte) ([]byte, error) {
	collectionKey, err := DecryptWithAlgorithm(encryptedCollectionKey.Ciphertext, encryptedCollectionKey.Nonce, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt collection key: %w", err)
	}

	return collectionKey, nil
}

// DecryptFileKey decrypts a file key using the collection key.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func DecryptFileKey(encryptedFileKey *EncryptedKey, collectionKey []byte) ([]byte, error) {
	fileKey, err := DecryptWithAlgorithm(encryptedFileKey.Ciphertext, encryptedFileKey.Nonce, collectionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt file key: %w", err)
	}

	return fileKey, nil
}

// DecryptPrivateKey decrypts the user's private key using the master key.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func DecryptPrivateKey(encryptedPrivateKey *EncryptedKey, masterKey []byte) ([]byte, error) {
	privateKey, err := DecryptWithAlgorithm(encryptedPrivateKey.Ciphertext, encryptedPrivateKey.Nonce, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt private key: %w", err)
	}

	return privateKey, nil
}

// DecryptRecoveryKey decrypts the user's recovery key using the master key.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func DecryptRecoveryKey(encryptedRecoveryKey *EncryptedKey, masterKey []byte) ([]byte, error) {
	recoveryKey, err := DecryptWithAlgorithm(encryptedRecoveryKey.Ciphertext, encryptedRecoveryKey.Nonce, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt recovery key: %w", err)
	}

	return recoveryKey, nil
}

// DecryptMasterKeyWithRecoveryKey decrypts the master key using the recovery key.
// This is used during account recovery.
// Auto-detects cipher based on nonce size (12 for ChaCha20, 24 for XSalsa20).
func DecryptMasterKeyWithRecoveryKey(encryptedMasterKey *EncryptedKey, recoveryKey []byte) ([]byte, error) {
	masterKey, err := DecryptWithAlgorithm(encryptedMasterKey.Ciphertext, encryptedMasterKey.Nonce, recoveryKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt master key with recovery key: %w", err)
	}

	return masterKey, nil
}

// GenerateMasterKey generates a new random master key.
func GenerateMasterKey() ([]byte, error) {
	return GenerateRandomBytes(MasterKeySize)
}

// GenerateCollectionKey generates a new random collection key.
func GenerateCollectionKey() ([]byte, error) {
	return GenerateRandomBytes(CollectionKeySize)
}

// GenerateFileKey generates a new random file key.
func GenerateFileKey() ([]byte, error) {
	return GenerateRandomBytes(FileKeySize)
}

// GenerateRecoveryKey generates a new random recovery key.
func GenerateRecoveryKey() ([]byte, error) {
	return GenerateRandomBytes(RecoveryKeySize)
}

// GenerateSalt generates a new random salt for password derivation.
func GenerateSalt() ([]byte, error) {
	return GenerateRandomBytes(Argon2SaltSize)
}

// EncryptMasterKey encrypts a master key with the KEK.
func (k *KeyChain) EncryptMasterKey(masterKey []byte) (*EncryptedKey, error) {
	if k.kek == nil {
		return nil, fmt.Errorf("keychain has been cleared")
	}

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

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

// EncryptCollectionKey encrypts a collection key with the master key using ChaCha20-Poly1305.
// For web frontend compatibility, use EncryptCollectionKeySecretBox instead.
func EncryptCollectionKey(collectionKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := Encrypt(collectionKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt collection key: %w", err)
	}

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

// EncryptCollectionKeySecretBox encrypts a collection key with the master key using XSalsa20-Poly1305.
// This is compatible with the web frontend's libsodium implementation.
func EncryptCollectionKeySecretBox(collectionKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := EncryptWithSecretBox(collectionKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt collection key: %w", err)
	}

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

// EncryptFileKey encrypts a file key with the collection key.
// NOTE: This uses ChaCha20-Poly1305 (12-byte nonce). For web frontend compatibility,
// use EncryptFileKeySecretBox instead.
func EncryptFileKey(fileKey, collectionKey []byte) (*EncryptedKey, error) {
	encrypted, err := Encrypt(fileKey, collectionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt file key: %w", err)
	}

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

// EncryptFileKeySecretBox encrypts a file key with the collection key using XSalsa20-Poly1305.
// This is compatible with the web frontend's libsodium implementation.
func EncryptFileKeySecretBox(fileKey, collectionKey []byte) (*EncryptedKey, error) {
	encrypted, err := EncryptWithSecretBox(fileKey, collectionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt file key: %w", err)
	}

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

// EncryptPrivateKey encrypts a private key with the master key.
func EncryptPrivateKey(privateKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := Encrypt(privateKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt private key: %w", err)
	}

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

// EncryptRecoveryKey encrypts a recovery key with the master key.
func EncryptRecoveryKey(recoveryKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := Encrypt(recoveryKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt recovery key: %w", err)
	}

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

// EncryptMasterKeyWithRecoveryKey encrypts a master key with the recovery key.
// This is used to enable account recovery.
func EncryptMasterKeyWithRecoveryKey(masterKey, recoveryKey []byte) (*EncryptedKey, error) {
	encrypted, err := Encrypt(masterKey, recoveryKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt master key with recovery key: %w", err)
	}

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

// =============================================================================
// SecretBox (XSalsa20-Poly1305) Encryption Functions
// These match the web frontend's libsodium crypto_secretbox_easy implementation
// =============================================================================

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

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

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

// EncryptPrivateKeySecretBox encrypts a private key with the master key using XSalsa20-Poly1305.
func EncryptPrivateKeySecretBox(privateKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := EncryptWithSecretBox(privateKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt private key: %w", err)
	}

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

// EncryptRecoveryKeySecretBox encrypts a recovery key with the master key using XSalsa20-Poly1305.
func EncryptRecoveryKeySecretBox(recoveryKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := EncryptWithSecretBox(recoveryKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt recovery key: %w", err)
	}

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

// EncryptMasterKeyWithRecoveryKeySecretBox encrypts a master key with the recovery key using XSalsa20-Poly1305.
func EncryptMasterKeyWithRecoveryKeySecretBox(masterKey, recoveryKey []byte) (*EncryptedKey, error) {
	encrypted, err := EncryptWithSecretBox(masterKey, recoveryKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt master key with recovery key: %w", err)
	}

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

// EncryptCollectionKeyForSharing encrypts a collection key for a recipient using BoxSeal.
// This is used when sharing a collection with another user.
func EncryptCollectionKeyForSharing(collectionKey, recipientPublicKey []byte) ([]byte, error) {
	if len(recipientPublicKey) != BoxPublicKeySize {
		return nil, fmt.Errorf("invalid recipient public key size: expected %d, got %d", BoxPublicKeySize, len(recipientPublicKey))
	}

	return EncryptWithBoxSeal(collectionKey, recipientPublicKey)
}

// DecryptSharedCollectionKey decrypts a collection key that was shared using BoxSeal.
// This is used when accessing a shared collection.
func DecryptSharedCollectionKey(encryptedCollectionKey, publicKey, privateKey []byte) ([]byte, error) {
	return DecryptWithBoxSeal(encryptedCollectionKey, publicKey, privateKey)
}

// ============================================================================
// Tag Key Operations
// ============================================================================

// GenerateTagKey generates a new 32-byte tag key for encrypting tag data.
func GenerateTagKey() ([]byte, error) {
	return GenerateRandomBytes(SecretBoxKeySize)
}

// GenerateKey is an alias for GenerateTagKey (convenience function).
func GenerateKey() []byte {
	key, _ := GenerateTagKey()
	return key
}

// EncryptTagKey encrypts a tag key with the master key using ChaCha20-Poly1305.
func EncryptTagKey(tagKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := Encrypt(tagKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt tag key: %w", err)
	}

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

// EncryptTagKeySecretBox encrypts a tag key with the master key using XSalsa20-Poly1305.
func EncryptTagKeySecretBox(tagKey, masterKey []byte) (*EncryptedKey, error) {
	encrypted, err := EncryptWithSecretBox(tagKey, masterKey)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt tag key: %w", err)
	}

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

// DecryptTagKey decrypts a tag key with the master key.
func DecryptTagKey(encryptedTagKey *EncryptedKey, masterKey []byte) ([]byte, error) {
	// Try XSalsa20-Poly1305 first (based on nonce size)
	if len(encryptedTagKey.Nonce) == SecretBoxNonceSize {
		return DecryptWithSecretBox(encryptedTagKey.Ciphertext, encryptedTagKey.Nonce, masterKey)
	}

	// Fall back to ChaCha20-Poly1305
	return Decrypt(encryptedTagKey.Ciphertext, encryptedTagKey.Nonce, masterKey)
}