package crypto

import (
	"crypto/sha256"
	"hash"

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

// HKDFInfo is the domain-separation string bound into the HKDF expansion.
// MUST match HKDF_INFO in crates/aura-crypto/src/kdf.rs.
var HKDFInfo = []byte("aura-v1-session")

// SessionKeyLen is the size of one directional AEAD key.
const SessionKeyLen = 32

// SessionKeys is the pair of directional 256-bit keys produced by the HKDF expansion.
type SessionKeys struct {
	ClientToServer [SessionKeyLen]byte
	ServerToClient [SessionKeyLen]byte
}

// DeriveSessionKeys runs HKDF-SHA256 with
//
//	salt = client_nonce || server_nonce (64 bytes)
//	IKM  = x25519_ss || mlkem_ss        (64 bytes)
//	info = "aura-v1-session", OKM 64 bytes -> (c2s, s2c)
//
// matching the production helper in crates/aura-crypto/src/kdf.rs byte-for-byte.
func DeriveSessionKeys(shared *HybridSharedSecret, clientNonce, serverNonce [32]byte) *SessionKeys {
	salt := make([]byte, 64)
	copy(salt[:32], clientNonce[:])
	copy(salt[32:], serverNonce[:])

	ikm := shared.Concat()
	hk := hkdf.New(func() hash.Hash { return sha256.New() }, ikm, salt, HKDFInfo)
	okm := make([]byte, 64)
	if _, err := hk.Read(okm); err != nil {
		// HKDF-Read for 64 bytes from SHA-256 is infallible; treat any error as a bug.
		panic(err)
	}
	var keys SessionKeys
	copy(keys.ClientToServer[:], okm[:32])
	copy(keys.ServerToClient[:], okm[32:])
	return &keys
}