package engine

import (
"crypto/ecdh"
"crypto/rand"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"

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

type KeyPair struct {
	PrivateKey *ecdh.PrivateKey
	PublicKey  *ecdh.PublicKey
}

type Engine struct{}

func NewEngine() *Engine {
	return &Engine{}
}

func (e *Engine) GenerateKeyPair() (*KeyPair, error) {
	curve := ecdh.X25519()
	priv, err := curve.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate private key: %w", err)
	}
	return &KeyPair{
		PrivateKey: priv,
		PublicKey:  priv.PublicKey(),
	}, nil
}

func (kp *KeyPair) EncodePublicKey() string {
	return hex.EncodeToString(kp.PublicKey.Bytes())
}

func (kp *KeyPair) EncodePrivateKey() string {
	return hex.EncodeToString(kp.PrivateKey.Bytes())
}

func (e *Engine) DecodePrivateKey(hexKey string) (*ecdh.PrivateKey, error) {
	bytes, err := hex.DecodeString(hexKey)
	if err != nil {
		return nil, err
	}
	return ecdh.X25519().NewPrivateKey(bytes)
}

func (e *Engine) DecodePublicKey(hexKey string) (*ecdh.PublicKey, error) {
	bytes, err := hex.DecodeString(hexKey)
	if err != nil {
		return nil, err
	}
	return ecdh.X25519().NewPublicKey(bytes)
}

// 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 := ecdh.X25519().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, chacha20poly1305.NonceSizeX)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Derive key from both secrets
	ikm := append(ss1, ss2...)
	kdf := hkdf.New(sha256.New, ikm, nonce, []byte("HIMBOCRYPT_V1"))
	symmetricKey := make([]byte, chacha20poly1305.KeySize)
	if _, err := io.ReadFull(kdf, symmetricKey); err != nil {
		return nil, fmt.Errorf("key derivation failed: %w", err)
	}

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

	// Pack it all up
	dst := make([]byte, 0, len(ephemeralPub.Bytes())+len(nonce)+len(plaintext)+aead.Overhead())
	dst = append(dst, ephemeralPub.Bytes()...)
	dst = append(dst, nonce...)
	return aead.Seal(dst, nonce, plaintext, nil), nil
}

func (e *Engine) Decrypt(recipientPriv *ecdh.PrivateKey, senderPub *ecdh.PublicKey, encryptedMsg []byte) ([]byte, error) {
	const pubKeySize = 32
	const nonceSize = 24

	if len(encryptedMsg) < pubKeySize+nonceSize {
		return nil, errors.New("message too short")
	}

	// Unpack
	ephemeralPubBytes := encryptedMsg[:pubKeySize]
	nonce := encryptedMsg[pubKeySize : pubKeySize+nonceSize]
	ciphertext := encryptedMsg[pubKeySize+nonceSize:]

	ephemeralPub, err := ecdh.X25519().NewPublicKey(ephemeralPubBytes)
	if err != nil {
		return nil, fmt.Errorf("invalid ephemeral public key: %w", err)
	}

	// Reconstruct secrets
	ss1, err := recipientPriv.ECDH(ephemeralPub)
	if err != nil {
		return nil, fmt.Errorf("ecdh ss1 failed: %w", err)
	}

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

	// Derive key
	ikm := append(ss1, ss2...)
	kdf := hkdf.New(sha256.New, ikm, nonce, []byte("HIMBOCRYPT_V1"))
	symmetricKey := make([]byte, chacha20poly1305.KeySize)
	if _, err := io.ReadFull(kdf, symmetricKey); err != nil {
		return nil, fmt.Errorf("key derivation failed: %w", err)
	}

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

	return aead.Open(nil, nonce, ciphertext, nil)
}