package pqid_test

import (
	"crypto/rand"
	"errors"
	"testing"

	"github.com/drjones/quantum-arcade/pkg/pqid"
)

func newKey(t *testing.T) (*pqid.PublicKey, *pqid.PrivateKey) {
	t.Helper()
	pub, priv, err := pqid.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	return pub, priv
}

func TestSignAndVerify(t *testing.T) {
	pub, priv := newKey(t)
	msg := []byte("authenticate me")

	sig, err := pqid.Sign(priv, msg)
	if err != nil {
		t.Fatal(err)
	}
	if err := pqid.Verify(pub, msg, sig); err != nil {
		t.Fatalf("valid hybrid signature rejected: %v", err)
	}
}

func TestWrongMessageFails(t *testing.T) {
	pub, priv := newKey(t)
	sig, _ := pqid.Sign(priv, []byte("original"))
	if err := pqid.Verify(pub, []byte("tampered"), sig); err == nil {
		t.Fatal("signature verified against a different message")
	}
}

func TestWrongKeyFails(t *testing.T) {
	_, priv := newKey(t)
	other, _ := newKey(t)
	sig, _ := pqid.Sign(priv, []byte("msg"))
	if err := pqid.Verify(other, []byte("msg"), sig); err == nil {
		t.Fatal("signature verified under the wrong key")
	}
}

// The whole point of hybrid: forging only the classical half must not
// authenticate. This is the quantum-adversary scenario — Shor breaks Ed25519,
// ML-DSA still holds.
func TestValidClassicalWithBrokenPostQuantumFails(t *testing.T) {
	pub, priv := newKey(t)
	msg := []byte("msg")
	sig, _ := pqid.Sign(priv, msg)

	// Keep the genuine Ed25519 signature, corrupt the ML-DSA half.
	forged := append([]byte(nil), sig...)
	forged[pqid.EdSignatureSize+10] ^= 0xff

	err := pqid.Verify(pub, msg, forged)
	if err == nil {
		t.Fatal("a broken post-quantum half still authenticated")
	}
	if !errors.Is(err, pqid.ErrPostQuantumFailed) {
		t.Fatalf("got %v, want ErrPostQuantumFailed", err)
	}
}

// And the mirror case: if lattice cryptography turns out to be weak, Ed25519
// must still stand in the way.
func TestValidPostQuantumWithBrokenClassicalFails(t *testing.T) {
	pub, priv := newKey(t)
	msg := []byte("msg")
	sig, _ := pqid.Sign(priv, msg)

	forged := append([]byte(nil), sig...)
	forged[5] ^= 0xff // corrupt the Ed25519 half

	err := pqid.Verify(pub, msg, forged)
	if err == nil {
		t.Fatal("a broken classical half still authenticated")
	}
	if !errors.Is(err, pqid.ErrClassicalFailed) {
		t.Fatalf("got %v, want ErrClassicalFailed", err)
	}
}

func TestSignatureSizeIsExact(t *testing.T) {
	_, priv := newKey(t)
	sig, err := pqid.Sign(priv, []byte("msg"))
	if err != nil {
		t.Fatal(err)
	}
	if len(sig) != pqid.SignatureSize {
		t.Fatalf("signature is %d bytes, want %d", len(sig), pqid.SignatureSize)
	}
}

func TestTruncatedSignatureRejected(t *testing.T) {
	pub, priv := newKey(t)
	sig, _ := pqid.Sign(priv, []byte("msg"))

	for _, n := range []int{0, 64, pqid.SignatureSize - 1} {
		if err := pqid.Verify(pub, []byte("msg"), sig[:n]); !errors.Is(err, pqid.ErrMalformedSignature) {
			t.Fatalf("signature truncated to %d bytes gave %v", n, err)
		}
	}
}

func TestPublicKeyRoundTrip(t *testing.T) {
	pub, priv := newKey(t)

	restored, err := pqid.ParsePublicKeyHex(pub.Hex())
	if err != nil {
		t.Fatal(err)
	}
	// The restored key must verify signatures made by the original.
	sig, _ := pqid.Sign(priv, []byte("msg"))
	if err := pqid.Verify(restored, []byte("msg"), sig); err != nil {
		t.Fatalf("round-tripped key failed to verify: %v", err)
	}
	if string(restored.ID()) != string(pub.ID()) {
		t.Fatal("round-tripped key has a different ID")
	}
}

func TestMalformedKeysRejected(t *testing.T) {
	cases := map[string][]byte{
		"empty":     {},
		"too short": make([]byte, pqid.PublicKeySize-1),
		"too long":  make([]byte, pqid.PublicKeySize+1),
	}
	for name, b := range cases {
		if _, err := pqid.ParsePublicKey(b); !errors.Is(err, pqid.ErrMalformedKey) {
			t.Errorf("%s: got %v, want ErrMalformedKey", name, err)
		}
	}
	if _, err := pqid.ParsePublicKeyHex("nothex!!"); !errors.Is(err, pqid.ErrMalformedKey) {
		t.Errorf("non-hex: got %v, want ErrMalformedKey", err)
	}
}

func TestIDIsStableAndDistinct(t *testing.T) {
	a, _ := newKey(t)
	b, _ := newKey(t)

	if string(a.ID()) != string(a.ID()) {
		t.Fatal("ID is not stable across calls")
	}
	if string(a.ID()) == string(b.ID()) {
		t.Fatal("two distinct keys produced the same ID")
	}
	if len(a.ID()) != pqid.IDSize {
		t.Fatalf("ID is %d bytes, want %d", len(a.ID()), pqid.IDSize)
	}
}

// Signatures must be bound to this application, so one captured from another
// ML-DSA protocol cannot be replayed here.
func TestSignaturesAreDomainSeparated(t *testing.T) {
	pub, priv := newKey(t)
	msg := []byte("msg")
	sig, _ := pqid.Sign(priv, msg)

	// Verifying with the correct context succeeds (covered above). Here we
	// confirm the context is actually in use by checking that a signature made
	// over the same message still fails if the ML-DSA half is swapped for one
	// generated under a different context.
	other := make([]byte, pqid.PQSignatureSize)
	if err := signWithContext(priv, msg, []byte("some-other-protocol"), other); err != nil {
		t.Fatal(err)
	}
	forged := append(append([]byte(nil), sig[:pqid.EdSignatureSize]...), other...)
	if err := pqid.Verify(pub, msg, forged); !errors.Is(err, pqid.ErrPostQuantumFailed) {
		t.Fatalf("signature from another context was accepted: %v", err)
	}
}

func BenchmarkSign(b *testing.B) {
	_, priv, _ := pqid.GenerateKey(rand.Reader)
	msg := []byte("benchmark message")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := pqid.Sign(priv, msg); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkVerify(b *testing.B) {
	pub, priv, _ := pqid.GenerateKey(rand.Reader)
	msg := []byte("benchmark message")
	sig, _ := pqid.Sign(priv, msg)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := pqid.Verify(pub, msg, sig); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkGenerateKey(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if _, _, err := pqid.GenerateKey(rand.Reader); err != nil {
			b.Fatal(err)
		}
	}
}

// A key stored by the browser and restored on the next visit must produce
// signatures the server still accepts.
func TestPrivateKeyRoundTripThroughStorage(t *testing.T) {
	pub, priv := newKey(t)

	// What the browser would persist.
	edSeed := priv.Ed.Seed()
	pqBytes, err := priv.PQ.MarshalBinary()
	if err != nil {
		t.Fatal(err)
	}

	restored, err := pqid.PrivateFromBytes(edSeed, pqBytes)
	if err != nil {
		t.Fatal(err)
	}

	msg := []byte("a challenge issued after the page reloaded")
	sig, err := pqid.Sign(restored, msg)
	if err != nil {
		t.Fatal(err)
	}
	if err := pqid.Verify(pub, msg, sig); err != nil {
		t.Fatalf("a restored key produced a signature the original public key rejects: %v", err)
	}
}

// The public key must be derivable, so a client cannot present one that does
// not match what it signs with.
func TestPublicKeyDerivesFromPrivate(t *testing.T) {
	pub, priv := newKey(t)

	derived, err := pqid.PublicFromPrivate(priv)
	if err != nil {
		t.Fatal(err)
	}
	if derived.Hex() != pub.Hex() {
		t.Fatal("derived public key does not match the generated one")
	}
	if string(derived.ID()) != string(pub.ID()) {
		t.Fatal("derived public key has a different account id")
	}
}

func TestMalformedStoredKeysRejected(t *testing.T) {
	_, priv := newKey(t)
	pqBytes, _ := priv.PQ.MarshalBinary()

	if _, err := pqid.PrivateFromBytes([]byte("short"), pqBytes); !errors.Is(err, pqid.ErrMalformedKey) {
		t.Errorf("short ed seed gave %v, want ErrMalformedKey", err)
	}
	if _, err := pqid.PrivateFromBytes(priv.Ed.Seed(), []byte("nonsense")); !errors.Is(err, pqid.ErrMalformedKey) {
		t.Errorf("bad pq key gave %v, want ErrMalformedKey", err)
	}
}