package crypto

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

	"filippo.io/age"
)

// Encryptor provides thread-safe encryption using age
type Encryptor struct {
	recipients []age.Recipient
	mu         sync.RWMutex
}

// NewEncryptor creates a new encryptor with the given age public keys
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
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
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
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 (thread-safe)
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
}