// Package handshake implements the client side of the Aura handshake state machine — a direct
// port of crates/aura-proto/src/handshake.rs::client_handshake.
//
// Order of messages (fixed by the Rust impl; see protocol.md §6.2):
//
//	1. C->S  ClientHello (plaintext): x25519_pub(32) || mlkem_ek(1184) || client_nonce(32)
//	2. S->C  ServerHello (plaintext): x25519_ephemeral(32) || mlkem_ct(1088) || server_nonce(32)
//	-- both sides derive the hybrid shared secret + directional SessionKeys --
//	3. S->C  ServerAuth (encrypted under s2c): u16(cert_der_len) || server_leaf_cert_der || sig
//	4. C->S  ClientAuth (encrypted under c2s): u16(cert_der_len) || client_leaf_cert_der || sig
//	5. C->S  Finished   (encrypted under c2s): HMAC-SHA256(key_c2s, transcript)
//	6. S->C  Finished   (encrypted under s2c): HMAC-SHA256(key_s2c, transcript)
//
// transcript = SHA-256(ClientHello_frame || ServerHello_frame), over the full serialized frames
// (header + payload) exactly as transmitted.
package handshake

import (
	"crypto/ecdsa"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"io"

	"github.com/aura/singbox-aura/aura/crypto"
	"github.com/aura/singbox-aura/aura/frame"
)

// ClientConfig is what the standalone CLI / sing-box outbound passes into Client.
//
// CAPEM, CertPEM, KeyPEM are PEM-encoded blobs (newlines, BEGIN/END lines and all). ServerName
// is the DNS name we expect to find in the server cert's SAN — must match the cert the server
// presents.
type ClientConfig struct {
	CAPEM      []byte
	CertPEM    []byte
	KeyPEM     []byte // PKCS#8 PEM, ECDSA P-256
	ServerName string
}

// Client runs the client side of the handshake to completion.
//
// On success it returns:
//   - DerivedKeys: the (c2s, s2c) session keys to seed the datagram codecs.
//   - PeerID:      the verified server name (the same string we passed in, on success).
//
// The caller wraps `r` / `w` over whatever transport is in use (the UDP reliability adapter
// for plain UDP; a TCP stream for the TCP fallback; a paired pipe in tests).
type Result struct {
	C2S        [32]byte
	S2C        [32]byte
	Transcript [32]byte
	PeerID     string
}

// Client drives the handshake state machine end-to-end.
func Client(r io.Reader, w io.Writer, cfg *ClientConfig) (*Result, error) {
	if cfg == nil {
		return nil, errors.New("aura/handshake: nil config")
	}

	// (1) Generate our hybrid keypair + nonce, send ClientHello.
	priv, pub, err := crypto.GenerateHybridKeypair()
	if err != nil {
		return nil, fmt.Errorf("hybrid keygen: %w", err)
	}
	var clientNonce [32]byte
	if _, err := rand.Read(clientNonce[:]); err != nil {
		return nil, fmt.Errorf("client nonce: %w", err)
	}

	chPayload := make([]byte, 0, crypto.X25519Len+crypto.MLKEMEKLen+32)
	chPayload = append(chPayload, pub.X25519[:]...)
	chPayload = append(chPayload, pub.MLKEM...)
	chPayload = append(chPayload, clientNonce[:]...)
	if len(chPayload) != crypto.X25519Len+crypto.MLKEMEKLen+32 {
		return nil, fmt.Errorf("client hello wrong size: %d", len(chPayload))
	}
	chHeader, err := frame.EncodeHeader(frame.MsgClientHello, len(chPayload))
	if err != nil {
		return nil, err
	}
	if err := frame.WriteFrame(w, frame.MsgClientHello, chPayload); err != nil {
		return nil, fmt.Errorf("write ClientHello: %w", err)
	}
	chWire := append(append([]byte{}, chHeader[:]...), chPayload...)

	// (2) Read ServerHello.
	sh, err := readExpect(r, frame.MsgServerHello)
	if err != nil {
		return nil, err
	}
	const expectSHLen = crypto.X25519Len + crypto.MLKEMCTLen + 32
	if len(sh.Payload) != expectSHLen {
		return nil, fmt.Errorf("ServerHello: wrong length %d (want %d)", len(sh.Payload), expectSHLen)
	}
	ct := &crypto.HybridCiphertext{MLKEMCT: append([]byte{}, sh.Payload[crypto.X25519Len:crypto.X25519Len+crypto.MLKEMCTLen]...)}
	copy(ct.X25519Eph[:], sh.Payload[:crypto.X25519Len])
	var serverNonce [32]byte
	copy(serverNonce[:], sh.Payload[crypto.X25519Len+crypto.MLKEMCTLen:])

	shared, err := priv.Decapsulate(ct)
	if err != nil {
		return nil, fmt.Errorf("decapsulate: %w", err)
	}
	keys := crypto.DeriveSessionKeys(shared, clientNonce, serverNonce)

	// transcript = SHA-256(client_hello_wire || server_hello_wire) over the bytes as transmitted.
	hash := sha256.New()
	hash.Write(chWire)
	hash.Write(sh.WireBytes())
	var transcript [32]byte
	copy(transcript[:], hash.Sum(nil))

	// Two AEAD sessions: client seals under c2s, opens under s2c. The counters continue across
	// the handshake/data boundary, so we must keep using the same instances.
	aeadC2S, err := crypto.NewAeadSession(keys.ClientToServer[:])
	if err != nil {
		return nil, err
	}
	aeadS2C, err := crypto.NewAeadSession(keys.ServerToClient[:])
	if err != nil {
		return nil, err
	}

	// (3) Server -> client ServerAuth (encrypted under s2c).
	serverAuth, err := openHandshakeMsg(r, frame.MsgServerAuth, aeadS2C)
	if err != nil {
		return nil, fmt.Errorf("ServerAuth: %w", err)
	}
	serverCertDER, serverSig, err := splitCertAndSig(serverAuth)
	if err != nil {
		return nil, err
	}
	if err := verifyServerCert(serverCertDER, cfg.CAPEM, cfg.ServerName); err != nil {
		return nil, fmt.Errorf("verify server cert: %w", err)
	}
	if err := verifySignature(serverCertDER, transcript[:], serverSig); err != nil {
		return nil, fmt.Errorf("verify server signature: %w", err)
	}

	// (4) Client -> server ClientAuth (encrypted under c2s).
	clientCertDER, err := pemCertToDER(cfg.CertPEM)
	if err != nil {
		return nil, fmt.Errorf("client cert: %w", err)
	}
	clientSig, err := signTranscript(cfg.KeyPEM, transcript[:])
	if err != nil {
		return nil, fmt.Errorf("sign transcript: %w", err)
	}
	clientAuth := buildCertAndSig(clientCertDER, clientSig)
	if err := sealHandshakeMsg(w, frame.MsgClientAuth, aeadC2S, clientAuth); err != nil {
		return nil, fmt.Errorf("write ClientAuth: %w", err)
	}

	// (5) Client -> server Finished (encrypted under c2s).
	clientFinished := hmacSHA256(keys.ClientToServer[:], transcript[:])
	if err := sealHandshakeMsg(w, frame.MsgFinished, aeadC2S, clientFinished); err != nil {
		return nil, fmt.Errorf("write client Finished: %w", err)
	}

	// (6) Server -> client Finished: verify against expected.
	serverFinished, err := openHandshakeMsg(r, frame.MsgFinished, aeadS2C)
	if err != nil {
		return nil, fmt.Errorf("server Finished: %w", err)
	}
	expectedServerFinished := hmacSHA256(keys.ServerToClient[:], transcript[:])
	if !hmac.Equal(serverFinished, expectedServerFinished) {
		return nil, errors.New("aura/handshake: server Finished MAC mismatch")
	}

	return &Result{
		C2S:        keys.ClientToServer,
		S2C:        keys.ServerToClient,
		Transcript: transcript,
		PeerID:     cfg.ServerName,
	}, nil
}

// readExpect reads one frame from r and demands it be of type want. An Alert is converted into
// a typed error.
func readExpect(r io.Reader, want frame.MsgType) (*frame.RawFrame, error) {
	rf, err := frame.ReadFrame(r)
	if err != nil {
		return nil, err
	}
	if rf.MsgType == frame.MsgAlert {
		code := byte(0)
		if len(rf.Payload) > 0 {
			code = rf.Payload[0]
		}
		return nil, fmt.Errorf("aura/handshake: peer alert code %d", code)
	}
	if rf.MsgType != want {
		return nil, fmt.Errorf("aura/handshake: expected %s, got %s", want, rf.MsgType)
	}
	return rf, nil
}

// sealHandshakeMsg seals plaintext under aead (advancing its counter) and writes one frame.
// AAD is the 5-byte header — same convention as Data records.
func sealHandshakeMsg(w io.Writer, msgType frame.MsgType, aead *crypto.AeadSession, plaintext []byte) error {
	sealedLen := len(plaintext) + 16 // Poly1305 tag
	hdr, err := frame.EncodeHeader(msgType, sealedLen)
	if err != nil {
		return err
	}
	ct := aead.Seal(plaintext, hdr[:])
	if len(ct) != sealedLen {
		return fmt.Errorf("aura/handshake: sealed wrong size %d (want %d)", len(ct), sealedLen)
	}
	return frame.WriteFrame(w, msgType, ct)
}

// openHandshakeMsg reads one frame of type msgType and AEAD-opens it.
func openHandshakeMsg(r io.Reader, msgType frame.MsgType, aead *crypto.AeadSession) ([]byte, error) {
	rf, err := readExpect(r, msgType)
	if err != nil {
		return nil, err
	}
	return aead.Open(rf.Payload, rf.Header[:])
}

// buildCertAndSig: u16_be(cert_der_len) || cert_der || signature.
func buildCertAndSig(certDER, sig []byte) []byte {
	out := make([]byte, 0, 2+len(certDER)+len(sig))
	var lb [2]byte
	binary.BigEndian.PutUint16(lb[:], uint16(len(certDER)))
	out = append(out, lb[:]...)
	out = append(out, certDER...)
	out = append(out, sig...)
	return out
}

// splitCertAndSig is the inverse.
func splitCertAndSig(buf []byte) (certDER, sig []byte, err error) {
	if len(buf) < 2 {
		return nil, nil, errors.New("aura/handshake: Auth: missing cert length")
	}
	certLen := int(binary.BigEndian.Uint16(buf[:2]))
	if len(buf) < 2+certLen {
		return nil, nil, errors.New("aura/handshake: Auth: truncated cert")
	}
	certDER = buf[2 : 2+certLen]
	sig = buf[2+certLen:]
	if len(sig) == 0 {
		return nil, nil, errors.New("aura/handshake: Auth: empty signature")
	}
	return certDER, sig, nil
}

// hmacSHA256 returns HMAC-SHA256(key, msg).
func hmacSHA256(key, msg []byte) []byte {
	m := hmac.New(sha256.New, key)
	m.Write(msg)
	return m.Sum(nil)
}

// pemCertToDER decodes the first CERTIFICATE PEM block.
func pemCertToDER(pemBytes []byte) ([]byte, error) {
	rest := pemBytes
	for {
		block, r := pem.Decode(rest)
		if block == nil {
			return nil, errors.New("aura/handshake: no CERTIFICATE block in PEM")
		}
		if block.Type == "CERTIFICATE" {
			return block.Bytes, nil
		}
		rest = r
	}
}

// pemKeyToDER decodes the first PRIVATE KEY-style PEM block. ECDSA leaves typically use PKCS#8
// ("PRIVATE KEY"); we also accept the old "EC PRIVATE KEY" form for compatibility.
func pemKeyToDER(pemBytes []byte) ([]byte, error) {
	rest := pemBytes
	for {
		block, r := pem.Decode(rest)
		if block == nil {
			return nil, errors.New("aura/handshake: no private-key block in PEM")
		}
		switch block.Type {
		case "PRIVATE KEY", "EC PRIVATE KEY", "RSA PRIVATE KEY":
			return block.Bytes, nil
		}
		rest = r
	}
}

// signTranscript signs a 32-byte transcript with the ECDSA P-256 PKCS#8 key in PEM form. The
// signature is the ASN.1 DER encoding ring uses on the Rust side (ECDSA_P256_SHA256_ASN1).
func signTranscript(keyPEM, transcript []byte) ([]byte, error) {
	if len(transcript) != 32 {
		return nil, fmt.Errorf("transcript must be 32 bytes, got %d", len(transcript))
	}
	der, err := pemKeyToDER(keyPEM)
	if err != nil {
		return nil, err
	}
	parsed, err := x509.ParsePKCS8PrivateKey(der)
	if err != nil {
		// Fall back to the old EC-specific encoding (rfc 5915).
		ec, err2 := x509.ParseECPrivateKey(der)
		if err2 != nil {
			return nil, fmt.Errorf("parse client key: pkcs8=%v ec=%v", err, err2)
		}
		parsed = ec
	}
	key, ok := parsed.(*ecdsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("aura/handshake: client key is %T, want *ecdsa.PrivateKey", parsed)
	}
	// ecdsa.SignASN1 returns the same ASN.1 DER (r,s) encoding ring produces.
	sig, err := ecdsa.SignASN1(rand.Reader, key, transcript)
	if err != nil {
		return nil, fmt.Errorf("ecdsa sign: %w", err)
	}
	return sig, nil
}

// verifySignature checks an ECDSA P-256/SHA-256 signature (ASN.1 DER) over the 32-byte transcript
// against the leaf cert's public key.
func verifySignature(certDER, transcript, sig []byte) error {
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return fmt.Errorf("parse peer cert: %w", err)
	}
	pub, ok := cert.PublicKey.(*ecdsa.PublicKey)
	if !ok {
		return fmt.Errorf("peer key is %T, want *ecdsa.PublicKey", cert.PublicKey)
	}
	if !ecdsa.VerifyASN1(pub, transcript, sig) {
		return errors.New("aura/handshake: signature did not verify")
	}
	return nil
}

// verifyServerCert validates the server leaf against the CA PEM and the expected DNS name.
func verifyServerCert(certDER, caPEM []byte, serverName string) error {
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(caPEM) {
		return errors.New("aura/handshake: CA PEM contains no certs")
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return fmt.Errorf("parse server cert: %w", err)
	}
	opts := x509.VerifyOptions{
		Roots:     pool,
		DNSName:   serverName,
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	if _, err := cert.Verify(opts); err != nil {
		return fmt.Errorf("verify chain: %w", err)
	}
	return nil
}