// Package crypto implements the Aura primitives the Go client side needs: hybrid X25519 +
// ML-KEM-768 KEM, HKDF-SHA256 session-key derivation, ChaCha20-Poly1305 AEAD using the same
// LE(u64)||[0;4] nonce scheme the Rust side uses, and the HMAC-SHA256 port-knock token.
//
// All exported sizes match the on-wire constants in crates/aura-crypto and aura-proto:
//
//	X25519 public / shared secret    32 bytes
//	ML-KEM-768 encapsulation key    1184 bytes
//	ML-KEM-768 ciphertext           1088 bytes
//	ML-KEM-768 shared secret          32 bytes
//
// We use crypto/mlkem (Go 1.24+ stdlib) for the post-quantum half. The Rust side uses the
// `ml_kem` 0.3 crate; both are FIPS 203 ML-KEM-768. The shared secrets agree byte-for-byte —
// asserted in crypto_test.go against the KAT vector emitted by `tools/export-kat`.
package crypto

import (
	"crypto/ecdh"
	"crypto/mlkem"
	"crypto/rand"
	"errors"
	"fmt"
)

// Sizes of the hybrid KEM building blocks, all in bytes.
const (
	X25519Len   = 32
	MLKEMEKLen  = 1184
	MLKEMCTLen  = 1088
	MLKEMSSLen  = 32
	HybridSSLen = X25519Len + MLKEMSSLen
)

// HybridPublicKey is the client's public half: a 32-byte X25519 public key plus a 1184-byte
// ML-KEM-768 encapsulation key.
type HybridPublicKey struct {
	X25519 [X25519Len]byte
	MLKEM  []byte // 1184 bytes
}

// HybridPrivateKey is the client's secret half. We hold the high-level keys so encapsulate /
// decapsulate are simple method calls.
type HybridPrivateKey struct {
	x25519Priv *ecdh.PrivateKey
	mlkemDk    *mlkem.DecapsulationKey768
}

// HybridCiphertext is the server's response: its ephemeral X25519 public key plus the ML-KEM
// ciphertext.
type HybridCiphertext struct {
	X25519Eph [X25519Len]byte
	MLKEMCT   []byte // 1088 bytes
}

// HybridSharedSecret is the 64-byte concatenation x25519_ss || kyber_ss.
type HybridSharedSecret struct {
	X25519SS [X25519Len]byte
	MLKEMSS  [MLKEMSSLen]byte
}

// Concat returns x25519_ss || mlkem_ss in one slice (the IKM HKDF consumes).
func (h *HybridSharedSecret) Concat() []byte {
	out := make([]byte, HybridSSLen)
	copy(out[:X25519Len], h.X25519SS[:])
	copy(out[X25519Len:], h.MLKEMSS[:])
	return out
}

// GenerateHybridKeypair produces a fresh client hybrid keypair using the OS RNG. Used by the
// standalone CLI; tests that need determinism instead call NewHybridPrivateFromSeeds or
// reconstruct from explicit bytes.
func GenerateHybridKeypair() (*HybridPrivateKey, *HybridPublicKey, error) {
	x, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("x25519 keygen: %w", err)
	}
	dk, err := mlkem.GenerateKey768()
	if err != nil {
		return nil, nil, fmt.Errorf("ml-kem keygen: %w", err)
	}
	return buildHybrid(x, dk)
}

// NewHybridPrivateFromBytes reconstructs a hybrid private key from raw 32-byte X25519 seed and
// the 64-byte ML-KEM seed (d || z). Mirrors the deterministic constructor the export-kat tool
// uses so the Go side can drive a handshake against the same KAT vector.
func NewHybridPrivateFromBytes(x25519Priv [X25519Len]byte, mlkemSeed [64]byte) (*HybridPrivateKey, *HybridPublicKey, error) {
	// x25519: NewPrivateKey requires a 32-byte scalar. Go enforces clamping inside the curve.
	x, err := ecdh.X25519().NewPrivateKey(x25519Priv[:])
	if err != nil {
		return nil, nil, fmt.Errorf("x25519 from bytes: %w", err)
	}
	dk, err := mlkem.NewDecapsulationKey768(mlkemSeed[:])
	if err != nil {
		return nil, nil, fmt.Errorf("ml-kem from seed: %w", err)
	}
	return buildHybrid(x, dk)
}

func buildHybrid(x *ecdh.PrivateKey, dk *mlkem.DecapsulationKey768) (*HybridPrivateKey, *HybridPublicKey, error) {
	priv := &HybridPrivateKey{x25519Priv: x, mlkemDk: dk}
	pub := &HybridPublicKey{MLKEM: dk.EncapsulationKey().Bytes()}
	if len(pub.MLKEM) != MLKEMEKLen {
		return nil, nil, fmt.Errorf("ml-kem ek wrong length: %d", len(pub.MLKEM))
	}
	xPub := x.PublicKey().Bytes()
	if len(xPub) != X25519Len {
		return nil, nil, fmt.Errorf("x25519 pub wrong length: %d", len(xPub))
	}
	copy(pub.X25519[:], xPub)
	return priv, pub, nil
}

// Decapsulate runs the client-side decapsulation: ECDH against the server's ephemeral X25519
// plus ML-KEM-768 decapsulation under the stored secret key.
func (h *HybridPrivateKey) Decapsulate(ct *HybridCiphertext) (*HybridSharedSecret, error) {
	if len(ct.MLKEMCT) != MLKEMCTLen {
		return nil, fmt.Errorf("ml-kem ct wrong length: %d", len(ct.MLKEMCT))
	}
	peerPub, err := ecdh.X25519().NewPublicKey(ct.X25519Eph[:])
	if err != nil {
		return nil, fmt.Errorf("x25519 peer pub: %w", err)
	}
	xss, err := h.x25519Priv.ECDH(peerPub)
	if err != nil {
		return nil, fmt.Errorf("x25519 ecdh: %w", err)
	}
	if len(xss) != X25519Len {
		return nil, fmt.Errorf("x25519 ss wrong length: %d", len(xss))
	}
	kss, err := h.mlkemDk.Decapsulate(ct.MLKEMCT)
	if err != nil {
		return nil, fmt.Errorf("ml-kem decapsulate: %w", err)
	}
	if len(kss) != MLKEMSSLen {
		return nil, fmt.Errorf("ml-kem ss wrong length: %d", len(kss))
	}
	out := &HybridSharedSecret{}
	copy(out.X25519SS[:], xss)
	copy(out.MLKEMSS[:], kss)
	return out, nil
}

// Encapsulate is the server side of the handshake. Provided here purely so a Go-side end-to-end
// test can drive both halves in-process. The standalone client never calls this.
func (p *HybridPublicKey) Encapsulate() (*HybridCiphertext, *HybridSharedSecret, error) {
	if len(p.MLKEM) != MLKEMEKLen {
		return nil, nil, errors.New("hybrid pub: invalid ml-kem ek length")
	}
	eph, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("x25519 eph keygen: %w", err)
	}
	peer, err := ecdh.X25519().NewPublicKey(p.X25519[:])
	if err != nil {
		return nil, nil, fmt.Errorf("x25519 peer: %w", err)
	}
	xss, err := eph.ECDH(peer)
	if err != nil {
		return nil, nil, fmt.Errorf("x25519 ecdh: %w", err)
	}
	ek, err := mlkem.NewEncapsulationKey768(p.MLKEM)
	if err != nil {
		return nil, nil, fmt.Errorf("ml-kem ek parse: %w", err)
	}
	kss, kct := ek.Encapsulate()

	ct := &HybridCiphertext{MLKEMCT: kct}
	copy(ct.X25519Eph[:], eph.PublicKey().Bytes())
	ss := &HybridSharedSecret{}
	copy(ss.X25519SS[:], xss)
	copy(ss.MLKEMSS[:], kss)
	return ct, ss, nil
}