package enclave

import (
	"encoding/json"
	"fmt"

	"enclave/internal/keybase"
)

// Enclave wraps a Keybase with encryption capabilities using WebAuthn PRF-derived keys.
type Enclave struct {
	keybase       *keybase.Keybase
	encryptionKey []byte
}

// Config holds enclave configuration options.
type Config struct {
	PRFOutput []byte
}

// New creates a new Enclave with the given PRF output for key derivation.
func New(prfOutput []byte) (*Enclave, error) {
	if len(prfOutput) == 0 {
		return nil, fmt.Errorf("enclave: PRF output required")
	}

	key, err := DeriveEncryptionKey(prfOutput)
	if err != nil {
		return nil, fmt.Errorf("enclave: key derivation failed: %w", err)
	}

	kb, err := keybase.Open()
	if err != nil {
		SecureZero(key)
		return nil, fmt.Errorf("enclave: failed to open keybase: %w", err)
	}

	return &Enclave{
		keybase:       kb,
		encryptionKey: key,
	}, nil
}

// Keybase returns the underlying keybase instance.
func (e *Enclave) Keybase() *keybase.Keybase {
	return e.keybase
}

// DID returns the current DID.
func (e *Enclave) DID() string {
	return e.keybase.DID()
}

// IsInitialized returns true if the enclave has been initialized with a DID.
func (e *Enclave) IsInitialized() bool {
	return e.keybase.IsInitialized()
}

// SerializeEncrypted exports the database state as encrypted bytes.
func (e *Enclave) SerializeEncrypted() ([]byte, error) {
	plaintext, err := e.keybase.Serialize()
	if err != nil {
		return nil, fmt.Errorf("enclave: serialization failed: %w", err)
	}

	encrypted, err := EncryptBytes(e.encryptionKey, plaintext)
	if err != nil {
		SecureZero(plaintext)
		return nil, fmt.Errorf("enclave: encryption failed: %w", err)
	}

	SecureZero(plaintext)
	return encrypted, nil
}

// LoadEncrypted loads the database state from encrypted bytes.
func (e *Enclave) LoadEncrypted(encryptedData []byte) error {
	plaintext, err := DecryptBytes(e.encryptionKey, encryptedData)
	if err != nil {
		return fmt.Errorf("enclave: decryption failed: %w", err)
	}
	defer SecureZero(plaintext)

	return e.loadFromPlaintext(plaintext)
}

// loadFromPlaintext parses and executes the SQL statements to restore database state.
func (e *Enclave) loadFromPlaintext(data []byte) error {
	if len(data) == 0 {
		return fmt.Errorf("enclave: empty data")
	}

	return e.keybase.RestoreFromDump(data)
}

// Close securely closes the enclave and zeros out sensitive data.
func (e *Enclave) Close() error {
	SecureZero(e.encryptionKey)
	return keybase.Close()
}

// EncryptedBundle represents a complete encrypted database export.
type EncryptedBundle struct {
	Version    int    `json:"version"`
	DID        string `json:"did"`
	Ciphertext []byte `json:"ciphertext"`
	Nonce      []byte `json:"nonce"`
}

// Export creates a complete encrypted bundle for storage.
func (e *Enclave) Export() (*EncryptedBundle, error) {
	plaintext, err := e.keybase.Serialize()
	if err != nil {
		return nil, fmt.Errorf("enclave: serialization failed: %w", err)
	}
	defer SecureZero(plaintext)

	encData, err := Encrypt(e.encryptionKey, plaintext)
	if err != nil {
		return nil, fmt.Errorf("enclave: encryption failed: %w", err)
	}

	return &EncryptedBundle{
		Version:    encData.Version,
		DID:        e.keybase.DID(),
		Ciphertext: encData.Ciphertext,
		Nonce:      encData.Nonce,
	}, nil
}

// Import loads an encrypted bundle.
func (e *Enclave) Import(bundle *EncryptedBundle) error {
	if bundle == nil {
		return fmt.Errorf("enclave: bundle cannot be nil")
	}

	encData := &EncryptedData{
		Version:    bundle.Version,
		Ciphertext: bundle.Ciphertext,
		Nonce:      bundle.Nonce,
	}

	plaintext, err := Decrypt(e.encryptionKey, encData)
	if err != nil {
		return fmt.Errorf("enclave: decryption failed: %w", err)
	}
	defer SecureZero(plaintext)

	return e.loadFromPlaintext(plaintext)
}

// MarshalBundle serializes an encrypted bundle to JSON.
func (b *EncryptedBundle) Marshal() ([]byte, error) {
	return json.Marshal(b)
}

// UnmarshalBundle deserializes an encrypted bundle from JSON.
func UnmarshalBundle(data []byte) (*EncryptedBundle, error) {
	var bundle EncryptedBundle
	if err := json.Unmarshal(data, &bundle); err != nil {
		return nil, fmt.Errorf("enclave: failed to unmarshal bundle: %w", err)
	}
	return &bundle, nil
}

// FromExisting wraps an existing keybase with encryption capabilities.
func FromExisting(kb *keybase.Keybase, prfOutput []byte) (*Enclave, error) {
	if kb == nil {
		return nil, fmt.Errorf("enclave: keybase cannot be nil")
	}
	if len(prfOutput) == 0 {
		return nil, fmt.Errorf("enclave: PRF output required")
	}

	key, err := DeriveEncryptionKey(prfOutput)
	if err != nil {
		return nil, fmt.Errorf("enclave: key derivation failed: %w", err)
	}

	return &Enclave{
		keybase:       kb,
		encryptionKey: key,
	}, nil
}
