package crypto

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"filippo.io/age"
)

// Encryptor provides thread-safe encryption using the age encryption library.
// It supports encrypting data for multiple recipients simultaneously, allowing
// any of the corresponding private keys to decrypt the data. This is useful
// for backup scenarios where multiple parties should be able to decrypt the data.
type Encryptor struct {
	recipients []age.Recipient
	mu         sync.RWMutex
}

// NewEncryptor creates a new encryptor with the given age public keys.
// Each public key should be a valid age X25519 recipient string (e.g., "age1...")
// At least one recipient must be provided. Returns an error if any of the
// public keys are invalid or if no recipients are specified.
func NewEncryptor(publicKeys []string) (*Encryptor, error) {
	if len(publicKeys) == 0 {
		return nil, fmt.Errorf("at least one recipient is required")
	}

	recipients := make([]age.Recipient, 0, len(publicKeys))
	for _, key := range publicKeys {
		recipient, err := age.ParseX25519Recipient(key)
		if err != nil {
			return nil, fmt.Errorf("parsing age recipient %s: %w", key, err)
		}
		recipients = append(recipients, recipient)
	}

	return &Encryptor{
		recipients: recipients,
	}, nil
}

// Encrypt encrypts data using age encryption for all configured recipients.
// The encrypted data can be decrypted by any of the corresponding private keys.
// This method is suitable for small to medium amounts of data that fit in memory.
// For large data streams, use EncryptStream or EncryptWriter instead.
func (e *Encryptor) Encrypt(data []byte) ([]byte, error) {
	e.mu.RLock()
	recipients := e.recipients
	e.mu.RUnlock()

	var buf bytes.Buffer

	// Create encrypted writer for all recipients
	w, err := age.Encrypt(&buf, recipients...)
	if err != nil {
		return nil, fmt.Errorf("creating encrypted writer: %w", err)
	}

	// Write data
	if _, err := w.Write(data); err != nil {
		return nil, fmt.Errorf("writing encrypted data: %w", err)
	}

	// Close to flush
	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("closing encrypted writer: %w", err)
	}

	return buf.Bytes(), nil
}

// EncryptStream encrypts data from reader to writer using age encryption.
// This method is suitable for encrypting large files or streams as it processes
// data in a streaming fashion without loading everything into memory.
// The encrypted data is written directly to the destination writer.
func (e *Encryptor) EncryptStream(dst io.Writer, src io.Reader) error {
	e.mu.RLock()
	recipients := e.recipients
	e.mu.RUnlock()

	// Create encrypted writer for all recipients
	w, err := age.Encrypt(dst, recipients...)
	if err != nil {
		return fmt.Errorf("creating encrypted writer: %w", err)
	}

	// Copy data
	if _, err := io.Copy(w, src); err != nil {
		return fmt.Errorf("copying encrypted data: %w", err)
	}

	// Close to flush
	if err := w.Close(); err != nil {
		return fmt.Errorf("closing encrypted writer: %w", err)
	}

	return nil
}

// EncryptWriter creates a writer that encrypts data written to it.
// All data written to the returned WriteCloser will be encrypted and written
// to the destination writer. The caller must call Close() on the returned
// writer to ensure all encrypted data is properly flushed and finalized.
// This is useful for integrating encryption into existing writer-based pipelines.
func (e *Encryptor) EncryptWriter(dst io.Writer) (io.WriteCloser, error) {
	e.mu.RLock()
	recipients := e.recipients
	e.mu.RUnlock()

	// Create encrypted writer for all recipients
	w, err := age.Encrypt(dst, recipients...)
	if err != nil {
		return nil, fmt.Errorf("creating encrypted writer: %w", err)
	}

	return w, nil
}

// UpdateRecipients updates the recipients for future encryption operations.
// This method is thread-safe and can be called while other encryption operations
// are in progress. Existing encryption operations will continue with the old
// recipients. At least one recipient must be provided. Returns an error if any
// of the public keys are invalid or if no recipients are specified.
func (e *Encryptor) UpdateRecipients(publicKeys []string) error {
	if len(publicKeys) == 0 {
		return fmt.Errorf("at least one recipient is required")
	}

	recipients := make([]age.Recipient, 0, len(publicKeys))
	for _, key := range publicKeys {
		recipient, err := age.ParseX25519Recipient(key)
		if err != nil {
			return fmt.Errorf("parsing age recipient %s: %w", key, err)
		}
		recipients = append(recipients, recipient)
	}

	e.mu.Lock()
	e.recipients = recipients
	e.mu.Unlock()

	return nil
}