// Package enclave provides encrypted database operations with WebAuthn PRF key derivation.
package enclave

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"io"

	"golang.org/x/crypto/hkdf"
)

const (
	// EnclaveSalt is the salt used for HKDF key derivation
	EnclaveSalt = "nebula-enclave-v1"

	// KeySize is the size of the derived encryption key (256 bits)
	KeySize = 32

	// NonceSize is the size of the GCM nonce (96 bits)
	NonceSize = 12

	// AuthTagSize is the size of the GCM authentication tag (128 bits)
	AuthTagSize = 16
)

// DeriveEncryptionKey derives a 256-bit encryption key from WebAuthn PRF output using HKDF.
//
// Parameters:
//   - prfOutput: The raw PRF output from WebAuthn (typically 32 bytes)
//
// Returns:
//   - A 32-byte key suitable for AES-256-GCM encryption
//   - An error if key derivation fails
func DeriveEncryptionKey(prfOutput []byte) ([]byte, error) {
	if len(prfOutput) == 0 {
		return nil, fmt.Errorf("enclave: PRF output cannot be empty")
	}

	salt := []byte(EnclaveSalt)
	info := []byte("database-encryption")

	hkdfReader := hkdf.New(sha256.New, prfOutput, salt, info)

	key := make([]byte, KeySize)
	if _, err := io.ReadFull(hkdfReader, key); err != nil {
		return nil, fmt.Errorf("enclave: failed to derive key: %w", err)
	}

	return key, nil
}

// DeriveKeyWithContext derives an encryption key with additional context binding.
// This allows deriving different keys for different purposes from the same PRF output.
//
// Parameters:
//   - prfOutput: The raw PRF output from WebAuthn
//   - context: Additional context to bind the key to (e.g., "database", "mpc-share")
func DeriveKeyWithContext(prfOutput []byte, context string) ([]byte, error) {
	if len(prfOutput) == 0 {
		return nil, fmt.Errorf("enclave: PRF output cannot be empty")
	}
	if context == "" {
		return nil, fmt.Errorf("enclave: context cannot be empty")
	}

	salt := []byte(EnclaveSalt)
	info := []byte(context)

	hkdfReader := hkdf.New(sha256.New, prfOutput, salt, info)

	key := make([]byte, KeySize)
	if _, err := io.ReadFull(hkdfReader, key); err != nil {
		return nil, fmt.Errorf("enclave: failed to derive key: %w", err)
	}

	return key, nil
}

// EncryptedData represents encrypted data with its metadata.
type EncryptedData struct {
	// Nonce is the unique nonce used for this encryption (12 bytes)
	Nonce []byte `json:"nonce"`
	// Ciphertext is the encrypted data including the GCM authentication tag
	Ciphertext []byte `json:"ciphertext"`
	// Version indicates the encryption scheme version
	Version int `json:"version"`
}

// Encrypt encrypts plaintext using AES-256-GCM with the provided key.
//
// Parameters:
//   - key: 32-byte encryption key (from DeriveEncryptionKey)
//   - plaintext: The data to encrypt
//
// Returns:
//   - EncryptedData containing nonce and ciphertext
//   - An error if encryption fails
func Encrypt(key, plaintext []byte) (*EncryptedData, error) {
	if len(key) != KeySize {
		return nil, fmt.Errorf("enclave: invalid key size: got %d, want %d", len(key), KeySize)
	}

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

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

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

	ciphertext := gcm.Seal(nil, nonce, plaintext, nil)

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

// Decrypt decrypts ciphertext using AES-256-GCM with the provided key.
//
// Parameters:
//   - key: 32-byte encryption key (from DeriveEncryptionKey)
//   - data: The EncryptedData to decrypt
//
// Returns:
//   - The decrypted plaintext
//   - An error if decryption fails (including authentication failure)
func Decrypt(key []byte, data *EncryptedData) ([]byte, error) {
	if len(key) != KeySize {
		return nil, fmt.Errorf("enclave: invalid key size: got %d, want %d", len(key), KeySize)
	}
	if data == nil {
		return nil, fmt.Errorf("enclave: encrypted data cannot be nil")
	}
	if len(data.Nonce) != NonceSize {
		return nil, fmt.Errorf("enclave: invalid nonce size: got %d, want %d", len(data.Nonce), NonceSize)
	}

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

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

	plaintext, err := gcm.Open(nil, data.Nonce, data.Ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("enclave: decryption failed (authentication error): %w", err)
	}

	return plaintext, nil
}

// EncryptBytes is a convenience function that encrypts and returns serialized bytes.
// The format is: version (1 byte) + nonce (12 bytes) + ciphertext (variable)
func EncryptBytes(key, plaintext []byte) ([]byte, error) {
	data, err := Encrypt(key, plaintext)
	if err != nil {
		return nil, err
	}

	result := make([]byte, 1+NonceSize+len(data.Ciphertext))
	result[0] = byte(data.Version)
	copy(result[1:1+NonceSize], data.Nonce)
	copy(result[1+NonceSize:], data.Ciphertext)

	return result, nil
}

// DecryptBytes is a convenience function that decrypts serialized encrypted bytes.
// Expected format: version (1 byte) + nonce (12 bytes) + ciphertext (variable)
func DecryptBytes(key, encryptedBytes []byte) ([]byte, error) {
	if len(encryptedBytes) < 1+NonceSize+AuthTagSize {
		return nil, fmt.Errorf("enclave: encrypted data too short")
	}

	version := int(encryptedBytes[0])
	if version != 1 {
		return nil, fmt.Errorf("enclave: unsupported encryption version: %d", version)
	}

	data := &EncryptedData{
		Version:    version,
		Nonce:      encryptedBytes[1 : 1+NonceSize],
		Ciphertext: encryptedBytes[1+NonceSize:],
	}

	return Decrypt(key, data)
}

// GenerateNonce generates a cryptographically secure random nonce.
func GenerateNonce() ([]byte, error) {
	nonce := make([]byte, NonceSize)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("enclave: failed to generate nonce: %w", err)
	}
	return nonce, nil
}

// SecureZero zeros out a byte slice to prevent sensitive data from remaining in memory.
func SecureZero(b []byte) {
	for i := range b {
		b[i] = 0
	}
}
