package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/hkdf"
	"crypto/sha256"
	"fmt"
)

// Hybrid encryption: X25519 (ephemeral sender key) + HKDF-SHA256 + AES-256-GCM.
//
// - Sender generates ephemeral X25519 keypair per message.
// - Derives aead key via HKDF using a random salt.
// - Uses random 12-byte iv for AES-GCM.

type HybridSealResult struct {
	EphemeralPublic []byte
	Salt            []byte
	IV              []byte
	Ciphertext      []byte
}

func ParseX25519PrivateKey(b64Priv string) (*ecdh.PrivateKey, error) {
	b, err := B64Decode(b64Priv)
	if err != nil {
		return nil, err
	}
	if len(b) != 32 {
		return nil, fmt.Errorf("invalid x25519 private key size: %d", len(b))
	}
	return ecdh.X25519().NewPrivateKey(b)
}

func ParseX25519PublicKey(b64Pub string) (*ecdh.PublicKey, error) {
	b, err := B64Decode(b64Pub)
	if err != nil {
		return nil, err
	}
	if len(b) != 32 {
		return nil, fmt.Errorf("invalid x25519 public key size: %d", len(b))
	}
	return ecdh.X25519().NewPublicKey(b)
}

func SealToX25519(recipientPub *ecdh.PublicKey, plaintext, aad []byte) (*HybridSealResult, error) {
	ePriv, err := ecdh.X25519().GenerateKey(randReader{})
	if err != nil {
		return nil, err
	}
	ss, err := ePriv.ECDH(recipientPub)
	if err != nil {
		return nil, err
	}
	salt, err := RandBytes(16)
	if err != nil {
		return nil, err
	}
	aeadKey, err := hkdfExpand(ss, salt, []byte("kong-agent-secure aead v1"), 32)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(aeadKey)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	iv, err := RandBytes(aead.NonceSize())
	if err != nil {
		return nil, err
	}
	ct := aead.Seal(nil, iv, plaintext, aad)
	return &HybridSealResult{
		EphemeralPublic: ePriv.PublicKey().Bytes(),
		Salt:            salt,
		IV:              iv,
		Ciphertext:      ct,
	}, nil
}

func OpenFromX25519(recipientPriv *ecdh.PrivateKey, epk, salt, iv, ciphertext, aad []byte) ([]byte, error) {
	ePub, err := ecdh.X25519().NewPublicKey(epk)
	if err != nil {
		return nil, err
	}
	ss, err := recipientPriv.ECDH(ePub)
	if err != nil {
		return nil, err
	}
	aeadKey, err := hkdfExpand(ss, salt, []byte("kong-agent-secure aead v1"), 32)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(aeadKey)
	if err != nil {
		return nil, err
	}
	aeadCipher, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(iv) != aeadCipher.NonceSize() {
		return nil, fmt.Errorf("invalid iv size: %d", len(iv))
	}
	pt, err := aeadCipher.Open(nil, iv, ciphertext, aad)
	if err != nil {
		return nil, err
	}
	return pt, nil
}

func hkdfExpand(secret, salt, info []byte, length int) ([]byte, error) {
	return hkdf.Key(sha256.New, secret, salt, string(info), length)
}

// randReader wraps crypto/rand.Reader without importing it here.
type randReader struct{}

func (randReader) Read(p []byte) (int, error) {
	b, err := RandBytes(len(p))
	if err != nil {
		return 0, err
	}
	copy(p, b)
	return len(p), nil
}