package handshake

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"testing"
	"time"

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

// TestSplitAndBuildCertAndSigRoundtrip: tiny but load-bearing — Auth payload layout must match
// the Rust wire format byte-for-byte.
func TestSplitAndBuildCertAndSigRoundtrip(t *testing.T) {
	cert := bytes.Repeat([]byte{0xAB}, 250)
	sig := []byte{0xCD, 0xEF, 0x01, 0x02}
	enc := buildCertAndSig(cert, sig)
	gotCert, gotSig, err := splitCertAndSig(enc)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(gotCert, cert) || !bytes.Equal(gotSig, sig) {
		t.Fatalf("roundtrip mismatch")
	}
	// Empty signature must be rejected.
	if _, _, err := splitCertAndSig(enc[:2+len(cert)]); err == nil {
		t.Fatal("empty sig must error")
	}
	// Truncated cert must be rejected.
	if _, _, err := splitCertAndSig(enc[:3]); err == nil {
		t.Fatal("truncated cert must error")
	}
}

// TestSignVerifyTranscriptRoundtrip: generate an ECDSA P-256 key + self-signed cert, sign a
// 32-byte transcript with our helper, verify with our helper, asserting we match the Rust side
// (ECDSA P-256 / SHA-256 / ASN.1 DER).
func TestSignVerifyTranscriptRoundtrip(t *testing.T) {
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	// Self-signed cert wrapping this key.
	tmpl := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "test-leaf"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(24 * time.Hour),
		BasicConstraintsValid: true,
		KeyUsage:              x509.KeyUsageDigitalSignature,
		SignatureAlgorithm:    x509.ECDSAWithSHA256,
	}
	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv)
	if err != nil {
		t.Fatal(err)
	}
	// Wrap our key in PKCS#8 PEM, as the production cert issuance does.
	keyDER, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER})

	var transcript [32]byte
	for i := range transcript {
		transcript[i] = byte(i ^ 0x55)
	}
	sig, err := signTranscript(keyPEM, transcript[:])
	if err != nil {
		t.Fatal(err)
	}
	if err := verifySignature(certDER, transcript[:], sig); err != nil {
		t.Fatal(err)
	}
	// Tampered transcript: verification must fail.
	bad := transcript
	bad[0] ^= 1
	if err := verifySignature(certDER, bad[:], sig); err == nil {
		t.Fatal("tampered transcript must fail")
	}
}

// TestClientHelloLayoutSize: sanity that we compute the expected hello payload size.
func TestClientHelloLayoutSize(t *testing.T) {
	const expected = 32 + 1184 + 32 // X25519 + ML-KEM ek + nonce
	if expected != 1248 {
		t.Fatalf("ClientHello expected size 1248, got %d", expected)
	}
	// And the on-wire frame adds the 5-byte header.
	hdr, err := frame.EncodeHeader(frame.MsgClientHello, expected)
	if err != nil {
		t.Fatal(err)
	}
	if hdr[0] != 0x01 || hdr[4] != 0x01 {
		t.Fatalf("header byte 0/4 mismatch: %x", hdr)
	}
}