package engine

import (
	"crypto/ecdh"
	"crypto/rand"
	"fmt"
	"io"

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

// Encrypt uses a hybrid scheme:
// 1. Ephemeral-Static ECDH (Forward Secrecy)
// 2. Static-Static ECDH (Auth)
// 3. XChaCha20-Poly1305 (AEAD)
// Output: [EphemeralPub] [Nonce] [Ciphertext]
func (e *Engine) Encrypt(senderPriv *ecdh.PrivateKey, recipientPub *ecdh.PublicKey, plaintext []byte) ([]byte, error) {
	// Generate ephemeral keys for forward secrecy
	ephemeralPriv, err := e.curve.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate ephemeral key: %w", err)
	}
	ephemeralPub := ephemeralPriv.PublicKey()

	// Calculate shared secrets
	ss1, err := ephemeralPriv.ECDH(recipientPub)
	if err != nil {
		return nil, fmt.Errorf("ecdh ss1 failed: %w", err)
	}

	ss2, err := senderPriv.ECDH(recipientPub)
	if err != nil {
		return nil, fmt.Errorf("ecdh ss2 failed: %w", err)
	}

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

	// Derive key from both secrets
	symmetricKey, err := deriveKey(ss1, ss2, nonce)
	if err != nil {
		return nil, err
	}

	aead, err := chacha20poly1305.NewX(symmetricKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create aead: %w", err)
	}

	// Pack it all up
	ephemeralPubBytes := ephemeralPub.Bytes()
	senderPubBytes := senderPriv.PublicKey().Bytes()

	// Bind ephemeral and sender public keys to ciphertext via AAD
	aad := buildAAD(ephemeralPubBytes, senderPubBytes)

	dst := make([]byte, 0, len(ephemeralPubBytes)+len(nonce)+len(plaintext)+aead.Overhead())
	dst = append(dst, ephemeralPubBytes...)
	dst = append(dst, nonce...)
	return aead.Seal(dst, nonce, plaintext, aad), nil
}