package keyexchange

import (
	"crypto/ecdh"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
)

var (
	// DefaultDHCurve is the default elliptic curve used for signing.
	DefaultDHCurve = ecdh.X25519
)

// GenerateDHKey generates a new ECDH private key for key exchange.
func GenerateDHKey() (*ecdh.PrivateKey, error) {
	return DefaultDHCurve().GenerateKey(rand.Reader)
}

// ComputeDHSecret computes the shared secret from the private key and the public key.
func ComputeDHSecret(priv *ecdh.PrivateKey, pub *ecdh.PublicKey) ([]byte, error) {
	secret, err := priv.ECDH(pub)
	if err != nil {
		return nil, err
	}

	key, err := hkdf.Key(sha256.New, secret, nil, "", 32)
	if err != nil {
		return nil, err
	}

	return key, nil
}

// ExportPublicKey exports the public key as a byte slice.
func ExportPublicKey(pub *ecdh.PublicKey) ([]byte, error) {
	return pub.Bytes(), nil
}

// ImportPublicKey imports the public key from a byte slice.
func ImportPublicKey(data []byte) (*ecdh.PublicKey, error) {
	return DefaultDHCurve().NewPublicKey(data)
}