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

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"io"

	"github.com/awnumar/memguard"
	"golang.org/x/crypto/argon2"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/nacl/secretbox"
	"golang.org/x/crypto/pbkdf2"
)

// KDF Algorithm identifiers
const (
	Argon2IDAlgorithm = "argon2id"
	PBKDF2Algorithm   = "PBKDF2-SHA256"
)

// Argon2id key derivation parameters
const (
	Argon2MemLimit    = 4 * 1024 * 1024 // 4 MB
	Argon2OpsLimit    = 1               // 1 iteration (time cost)
	Argon2Parallelism = 1               // 1 thread
	Argon2KeySize     = 32              // 256-bit output
	Argon2SaltSize    = 16              // 128-bit salt
)

// PBKDF2 key derivation parameters (matching web frontend)
const (
	PBKDF2Iterations = 100000 // 100,000 iterations (matching web frontend)
	PBKDF2KeySize    = 32     // 256-bit output
	PBKDF2SaltSize   = 16     // 128-bit salt
)

// ChaCha20-Poly1305 constants (IETF variant - 12 byte nonce)
const (
	ChaCha20Poly1305KeySize   = 32 // ChaCha20 key size
	ChaCha20Poly1305NonceSize = 12 // ChaCha20-Poly1305 nonce size
	ChaCha20Poly1305Overhead  = 16 // Poly1305 authentication tag size
)

// XSalsa20-Poly1305 (NaCl secretbox) constants - 24 byte nonce
// Used by web frontend (libsodium crypto_secretbox_easy)
const (
	SecretBoxKeySize   = 32 // Same as ChaCha20
	SecretBoxNonceSize = 24 // XSalsa20 uses 24-byte nonce
	SecretBoxOverhead  = secretbox.Overhead // 16 bytes (Poly1305 tag)
)

// Key sizes
const (
	MasterKeySize     = 32
	CollectionKeySize = 32
	FileKeySize       = 32
	RecoveryKeySize   = 32
)

// NaCl Box constants
const (
	BoxPublicKeySize = 32
	BoxSecretKeySize = 32
	BoxNonceSize     = 24
)

// EncryptedData represents encrypted data with its nonce.
type EncryptedData struct {
	Ciphertext []byte
	Nonce      []byte
}

// DeriveKeyFromPassword derives a key encryption key (KEK) from a password using Argon2id.
// This is the legacy function - prefer DeriveKeyFromPasswordWithAlgorithm for new code.
func DeriveKeyFromPassword(password string, salt []byte) ([]byte, error) {
	return DeriveKeyFromPasswordArgon2id(password, salt)
}

// DeriveKeyFromPasswordArgon2id derives a KEK using Argon2id algorithm.
// SECURITY: Password bytes are wiped from memory after key derivation.
func DeriveKeyFromPasswordArgon2id(password string, salt []byte) ([]byte, error) {
	if len(salt) != Argon2SaltSize {
		return nil, fmt.Errorf("invalid salt size: expected %d, got %d", Argon2SaltSize, len(salt))
	}

	passwordBytes := []byte(password)
	defer memguard.WipeBytes(passwordBytes) // SECURITY: Wipe password bytes after use

	key := argon2.IDKey(
		passwordBytes,
		salt,
		Argon2OpsLimit,    // time cost = 1
		Argon2MemLimit,    // memory = 4 MB
		Argon2Parallelism, // parallelism = 1
		Argon2KeySize,     // output size = 32 bytes
	)

	return key, nil
}

// DeriveKeyFromPasswordPBKDF2 derives a KEK using PBKDF2-SHA256 algorithm.
// This matches the web frontend's implementation.
// SECURITY: Password bytes are wiped from memory after key derivation.
func DeriveKeyFromPasswordPBKDF2(password string, salt []byte) ([]byte, error) {
	if len(salt) != PBKDF2SaltSize {
		return nil, fmt.Errorf("invalid salt size: expected %d, got %d", PBKDF2SaltSize, len(salt))
	}

	passwordBytes := []byte(password)
	defer memguard.WipeBytes(passwordBytes) // SECURITY: Wipe password bytes after use

	key := pbkdf2.Key(
		passwordBytes,
		salt,
		PBKDF2Iterations, // 100,000 iterations
		PBKDF2KeySize,    // 32 bytes output
		sha256.New,       // SHA-256 hash
	)

	return key, nil
}

// DeriveKeyFromPasswordWithAlgorithm derives a KEK using the specified algorithm.
// algorithm should be one of: Argon2IDAlgorithm, PBKDF2Algorithm
func DeriveKeyFromPasswordWithAlgorithm(password string, salt []byte, algorithm string) ([]byte, error) {
	switch algorithm {
	case Argon2IDAlgorithm: // "argon2id"
		return DeriveKeyFromPasswordArgon2id(password, salt)
	case PBKDF2Algorithm, "pbkdf2", "pbkdf2-sha256":
		return DeriveKeyFromPasswordPBKDF2(password, salt)
	default:
		return nil, fmt.Errorf("unsupported KDF algorithm: %s", algorithm)
	}
}

// Encrypt encrypts data with a symmetric key using ChaCha20-Poly1305.
func Encrypt(data, key []byte) (*EncryptedData, error) {
	if len(key) != ChaCha20Poly1305KeySize {
		return nil, fmt.Errorf("invalid key size: expected %d, got %d", ChaCha20Poly1305KeySize, len(key))
	}

	// Create ChaCha20-Poly1305 cipher
	cipher, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	// Generate random nonce (12 bytes for ChaCha20-Poly1305)
	nonce, err := GenerateRandomBytes(ChaCha20Poly1305NonceSize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Encrypt
	ciphertext := cipher.Seal(nil, nonce, data, nil)

	return &EncryptedData{
		Ciphertext: ciphertext,
		Nonce:      nonce,
	}, nil
}

// Decrypt decrypts data with a symmetric key using ChaCha20-Poly1305.
func Decrypt(ciphertext, nonce, key []byte) ([]byte, error) {
	if len(key) != ChaCha20Poly1305KeySize {
		return nil, fmt.Errorf("invalid key size: expected %d, got %d", ChaCha20Poly1305KeySize, len(key))
	}

	if len(nonce) != ChaCha20Poly1305NonceSize {
		return nil, fmt.Errorf("invalid nonce size: expected %d, got %d", ChaCha20Poly1305NonceSize, len(nonce))
	}

	// Create ChaCha20-Poly1305 cipher
	cipher, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	// Decrypt
	plaintext, err := cipher.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt: %w", err)
	}

	return plaintext, nil
}

// EncryptWithSecretBox encrypts data with a symmetric key using XSalsa20-Poly1305 (NaCl secretbox).
// This is compatible with libsodium's crypto_secretbox_easy used by the web frontend.
// SECURITY: Key arrays are wiped from memory after encryption.
func EncryptWithSecretBox(data, key []byte) (*EncryptedData, error) {
	if len(key) != SecretBoxKeySize {
		return nil, fmt.Errorf("invalid key size: expected %d, got %d", SecretBoxKeySize, len(key))
	}

	// Generate random nonce (24 bytes for XSalsa20)
	nonce, err := GenerateRandomBytes(SecretBoxNonceSize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Convert to fixed-size arrays for NaCl
	var keyArray [32]byte
	var nonceArray [24]byte
	copy(keyArray[:], key)
	copy(nonceArray[:], nonce)
	defer memguard.WipeBytes(keyArray[:]) // SECURITY: Wipe key array

	// Encrypt using secretbox
	ciphertext := secretbox.Seal(nil, data, &nonceArray, &keyArray)

	return &EncryptedData{
		Ciphertext: ciphertext,
		Nonce:      nonce,
	}, nil
}

// DecryptWithSecretBox decrypts data with a symmetric key using XSalsa20-Poly1305 (NaCl secretbox).
// This is compatible with libsodium's crypto_secretbox_open_easy used by the web frontend.
// SECURITY: Key arrays are wiped from memory after decryption.
func DecryptWithSecretBox(ciphertext, nonce, key []byte) ([]byte, error) {
	if len(key) != SecretBoxKeySize {
		return nil, fmt.Errorf("invalid key size: expected %d, got %d", SecretBoxKeySize, len(key))
	}

	if len(nonce) != SecretBoxNonceSize {
		return nil, fmt.Errorf("invalid nonce size: expected %d, got %d", SecretBoxNonceSize, len(nonce))
	}

	// Convert to fixed-size arrays for NaCl
	var keyArray [32]byte
	var nonceArray [24]byte
	copy(keyArray[:], key)
	copy(nonceArray[:], nonce)
	defer memguard.WipeBytes(keyArray[:]) // SECURITY: Wipe key array

	// Decrypt using secretbox
	plaintext, ok := secretbox.Open(nil, ciphertext, &nonceArray, &keyArray)
	if !ok {
		return nil, errors.New("failed to decrypt: invalid key, nonce, or corrupted ciphertext")
	}

	return plaintext, nil
}

// DecryptWithAlgorithm decrypts data using the appropriate cipher based on nonce size.
// - 12-byte nonce: ChaCha20-Poly1305 (IETF variant)
// - 24-byte nonce: XSalsa20-Poly1305 (NaCl secretbox)
func DecryptWithAlgorithm(ciphertext, nonce, key []byte) ([]byte, error) {
	switch len(nonce) {
	case ChaCha20Poly1305NonceSize: // 12 bytes
		return Decrypt(ciphertext, nonce, key)
	case SecretBoxNonceSize: // 24 bytes
		return DecryptWithSecretBox(ciphertext, nonce, key)
	default:
		return nil, fmt.Errorf("invalid nonce size: %d (expected %d for ChaCha20 or %d for XSalsa20)",
			len(nonce), ChaCha20Poly1305NonceSize, SecretBoxNonceSize)
	}
}

// EncryptWithBoxSeal encrypts data anonymously using NaCl sealed box.
// The result format is: ephemeral_public_key (32) || nonce (24) || ciphertext + auth_tag.
func EncryptWithBoxSeal(message []byte, recipientPublicKey []byte) ([]byte, error) {
	if len(recipientPublicKey) != BoxPublicKeySize {
		return nil, fmt.Errorf("recipient public key must be %d bytes", BoxPublicKeySize)
	}

	var recipientPubKey [32]byte
	copy(recipientPubKey[:], recipientPublicKey)

	// Generate ephemeral keypair
	ephemeralPubKey, ephemeralPrivKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate ephemeral keypair: %w", err)
	}

	// Generate random nonce
	nonce, err := GenerateRandomBytes(BoxNonceSize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}
	var nonceArray [24]byte
	copy(nonceArray[:], nonce)

	// Encrypt with ephemeral private key
	ciphertext := box.Seal(nil, message, &nonceArray, &recipientPubKey, ephemeralPrivKey)

	// Result format: ephemeral_public_key || nonce || ciphertext
	result := make([]byte, BoxPublicKeySize+BoxNonceSize+len(ciphertext))
	copy(result[:BoxPublicKeySize], ephemeralPubKey[:])
	copy(result[BoxPublicKeySize:BoxPublicKeySize+BoxNonceSize], nonce)
	copy(result[BoxPublicKeySize+BoxNonceSize:], ciphertext)

	return result, nil
}

// DecryptWithBoxSeal decrypts data that was encrypted with EncryptWithBoxSeal.
// SECURITY: Key arrays are wiped from memory after decryption.
func DecryptWithBoxSeal(sealedData []byte, recipientPublicKey, recipientPrivateKey []byte) ([]byte, error) {
	if len(recipientPublicKey) != BoxPublicKeySize {
		return nil, fmt.Errorf("recipient public key must be %d bytes", BoxPublicKeySize)
	}
	if len(recipientPrivateKey) != BoxSecretKeySize {
		return nil, fmt.Errorf("recipient private key must be %d bytes", BoxSecretKeySize)
	}
	if len(sealedData) < BoxPublicKeySize+BoxNonceSize+box.Overhead {
		return nil, errors.New("sealed data too short")
	}

	// Extract components
	ephemeralPublicKey := sealedData[:BoxPublicKeySize]
	nonce := sealedData[BoxPublicKeySize : BoxPublicKeySize+BoxNonceSize]
	ciphertext := sealedData[BoxPublicKeySize+BoxNonceSize:]

	// Create fixed-size arrays
	var ephemeralPubKey [32]byte
	var recipientPrivKey [32]byte
	var nonceArray [24]byte
	copy(ephemeralPubKey[:], ephemeralPublicKey)
	copy(recipientPrivKey[:], recipientPrivateKey)
	copy(nonceArray[:], nonce)
	defer memguard.WipeBytes(recipientPrivKey[:]) // SECURITY: Wipe private key array

	// Decrypt
	plaintext, ok := box.Open(nil, ciphertext, &nonceArray, &ephemeralPubKey, &recipientPrivKey)
	if !ok {
		return nil, errors.New("failed to decrypt sealed box: invalid keys or corrupted ciphertext")
	}

	return plaintext, nil
}

// DecryptAnonymousBox decrypts sealed box data (used in login challenges).
// SECURITY: Key arrays are wiped from memory after decryption.
func DecryptAnonymousBox(encryptedData []byte, recipientPublicKey, recipientPrivateKey []byte) ([]byte, error) {
	if len(recipientPublicKey) != BoxPublicKeySize {
		return nil, fmt.Errorf("recipient public key must be %d bytes", BoxPublicKeySize)
	}
	if len(recipientPrivateKey) != BoxSecretKeySize {
		return nil, fmt.Errorf("recipient private key must be %d bytes", BoxSecretKeySize)
	}

	var pubKeyArray, privKeyArray [32]byte
	copy(pubKeyArray[:], recipientPublicKey)
	copy(privKeyArray[:], recipientPrivateKey)
	defer memguard.WipeBytes(privKeyArray[:]) // SECURITY: Wipe private key array

	decryptedData, ok := box.OpenAnonymous(nil, encryptedData, &pubKeyArray, &privKeyArray)
	if !ok {
		return nil, errors.New("failed to decrypt anonymous box: invalid keys or corrupted data")
	}

	return decryptedData, nil
}

// GenerateRandomBytes generates cryptographically secure random bytes.
func GenerateRandomBytes(size int) ([]byte, error) {
	if size <= 0 {
		return nil, errors.New("size must be positive")
	}

	buf := make([]byte, size)
	_, err := io.ReadFull(rand.Reader, buf)
	if err != nil {
		return nil, fmt.Errorf("failed to generate random bytes: %w", err)
	}
	return buf, nil
}

// GenerateKeyPair generates a NaCl box keypair for asymmetric encryption.
func GenerateKeyPair() (publicKey []byte, privateKey []byte, err error) {
	pubKey, privKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate key pair: %w", err)
	}

	return pubKey[:], privKey[:], nil
}

// ClearBytes overwrites a byte slice with zeros using memguard for secure wiping.
// This should be called on sensitive data like keys when they're no longer needed.
// SECURITY: Uses memguard.WipeBytes for secure memory wiping that prevents compiler optimizations.
func ClearBytes(b []byte) {
	memguard.WipeBytes(b)
}

// CombineNonceAndCiphertext combines nonce and ciphertext into a single byte slice.
func CombineNonceAndCiphertext(nonce, ciphertext []byte) []byte {
	combined := make([]byte, len(nonce)+len(ciphertext))
	copy(combined[:len(nonce)], nonce)
	copy(combined[len(nonce):], ciphertext)
	return combined
}

// SplitNonceAndCiphertext splits a combined byte slice into nonce and ciphertext.
// This function defaults to ChaCha20-Poly1305 nonce size (12 bytes) for backward compatibility.
// For XSalsa20-Poly1305 (24-byte nonce), use SplitNonceAndCiphertextSecretBox.
func SplitNonceAndCiphertext(combined []byte) (nonce []byte, ciphertext []byte, err error) {
	if len(combined) < ChaCha20Poly1305NonceSize {
		return nil, nil, fmt.Errorf("combined data too short: expected at least %d bytes, got %d", ChaCha20Poly1305NonceSize, len(combined))
	}

	nonce = combined[:ChaCha20Poly1305NonceSize]
	ciphertext = combined[ChaCha20Poly1305NonceSize:]
	return nonce, ciphertext, nil
}

// SplitNonceAndCiphertextSecretBox splits a combined byte slice for XSalsa20-Poly1305 (24-byte nonce).
// This is compatible with libsodium's secretbox format: nonce (24) || ciphertext || mac (16).
func SplitNonceAndCiphertextSecretBox(combined []byte) (nonce []byte, ciphertext []byte, err error) {
	if len(combined) < SecretBoxNonceSize {
		return nil, nil, fmt.Errorf("combined data too short: expected at least %d bytes, got %d", SecretBoxNonceSize, len(combined))
	}

	nonce = combined[:SecretBoxNonceSize]
	ciphertext = combined[SecretBoxNonceSize:]
	return nonce, ciphertext, nil
}

// SplitNonceAndCiphertextAuto automatically detects the nonce size based on data length.
// It uses heuristics to determine if data is ChaCha20-Poly1305 (12-byte nonce) or XSalsa20 (24-byte nonce).
// This function should be used when the cipher type is unknown.
func SplitNonceAndCiphertextAuto(combined []byte) (nonce []byte, ciphertext []byte, err error) {
	// Web frontend uses XSalsa20-Poly1305 with 24-byte nonce
	// Native app used to use ChaCha20-Poly1305 with 12-byte nonce
	//
	// For encrypted master key data:
	// - Web frontend: nonce (24) + ciphertext (32 + 16 MAC) = 72 bytes
	// - Native/old:   nonce (12) + ciphertext (32 + 16 MAC) = 60 bytes
	//
	// We can distinguish by checking if the data length suggests 24-byte nonce
	// Data encrypted with 24-byte nonce will be 12 bytes longer than 12-byte nonce version

	if len(combined) < ChaCha20Poly1305NonceSize+ChaCha20Poly1305Overhead {
		return nil, nil, fmt.Errorf("combined data too short: expected at least %d bytes, got %d",
			ChaCha20Poly1305NonceSize+ChaCha20Poly1305Overhead, len(combined))
	}

	// If data length is at least 72 bytes (24 nonce + 32 key + 16 MAC for master key),
	// try XSalsa20 format first. This is the web frontend format.
	if len(combined) >= SecretBoxNonceSize+SecretBoxOverhead+1 {
		return SplitNonceAndCiphertextSecretBox(combined)
	}

	// Default to ChaCha20-Poly1305 (legacy)
	return SplitNonceAndCiphertext(combined)
}

// EncodeToBase64 encodes bytes to base64 standard encoding.
func EncodeToBase64(data []byte) string {
	return base64.StdEncoding.EncodeToString(data)
}

// DecodeFromBase64 decodes a base64 standard encoded string to bytes.
func DecodeFromBase64(s string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(s)
}