package gnarl

import (
	"math/big"
	"testing"
)

func TestPrimeProperties(t *testing.T) {
	p, q, n := Params()

	// P should be 216 bits.
	if p.BitLen() != 216 {
		t.Fatalf("P bits = %d, want 216", p.BitLen())
	}

	// Q should be ~213 bits.
	if q.BitLen() < 212 || q.BitLen() > 214 {
		t.Fatalf("Q bits = %d, want ~213", q.BitLen())
	}

	// P should be prime.
	if !p.ProbablyPrime(20) {
		t.Fatal("P is not prime")
	}

	// Q should be prime.
	if !q.ProbablyPrime(20) {
		t.Fatal("Q is not prime")
	}

	// P ≡ 2 mod 3.
	if new(big.Int).Mod(p, big.NewInt(3)).Int64() != 2 {
		t.Fatalf("P mod 3 = %d, want 2", new(big.Int).Mod(p, big.NewInt(3)).Int64())
	}

	// P mod 5 ∈ {2, 3}.
	pm5 := new(big.Int).Mod(p, big.NewInt(5)).Int64()
	if pm5 != 2 && pm5 != 3 {
		t.Fatalf("P mod 5 = %d, want 2 or 3", pm5)
	}

	// N = P + 1.
	if new(big.Int).Add(p, big.NewInt(1)).Cmp(n) != 0 {
		t.Fatal("N != P + 1")
	}

	// N = 6Q.
	sixQ := new(big.Int).Mul(big.NewInt(6), q)
	if n.Cmp(sixQ) != 0 {
		t.Fatal("N != 6Q")
	}

	// 5 is QNR mod P.
	pm1 := new(big.Int).Sub(p, big.NewInt(1))
	exp := new(big.Int).Rsh(pm1, 1)
	euler := new(big.Int).Exp(big.NewInt(5), exp, p)
	if euler.Cmp(pm1) != 0 {
		t.Fatal("5 is not QNR mod P")
	}
}

func TestDeterministicDerivation(t *testing.T) {
	// Running Params twice should give the same values (it's deterministic from seed).
	p1, q1, _ := Params()
	p2, q2, _ := Params()
	if p1.Cmp(p2) != 0 {
		t.Fatal("P is not deterministic")
	}
	if q1.Cmp(q2) != 0 {
		t.Fatal("Q is not deterministic")
	}
}

func TestFieldArithmeticIdentity(t *testing.T) {
	// 1 * 1 = 1 in Montgomery form.
	var r fe
	montMul(&r, &feOne, &feOne)
	if feEqual(&r, &feOne) != 1 {
		t.Fatal("1 * 1 != 1")
	}

	// 0 + 1 = 1.
	feAdd(&r, &feZero, &feOne)
	if feEqual(&r, &feOne) != 1 {
		t.Fatal("0 + 1 != 1")
	}

	// 1 - 1 = 0.
	feSub(&r, &feOne, &feOne)
	if feIsZero(&r) != 1 {
		t.Fatal("1 - 1 != 0")
	}

	// -0 = 0.
	feNeg(&r, &feZero)
	if feIsZero(&r) != 1 {
		t.Fatal("-0 != 0")
	}
}

func TestMontgomeryRoundTrip(t *testing.T) {
	// Convert 42 to Montgomery form and back.
	var a, b fe
	a = fe{42, 0, 0, 0}
	feToMont(&b, &a)
	feFromMont(&a, &b)
	if a[0] != 42 || a[1] != 0 || a[2] != 0 || a[3] != 0 {
		t.Fatalf("round-trip failed: got %v", a)
	}
}

func TestFieldMulCommutativity(t *testing.T) {
	var a, b, ab, ba fe
	feFromSmall(&a, 7)
	feFromSmall(&b, 13)
	montMul(&ab, &a, &b)
	montMul(&ba, &b, &a)
	if feEqual(&ab, &ba) != 1 {
		t.Fatal("a*b != b*a")
	}

	// Check that 7 * 13 = 91.
	var expected fe
	feFromSmall(&expected, 91)
	if feEqual(&ab, &expected) != 1 {
		t.Fatal("7 * 13 != 91")
	}
}

func TestFieldInverse(t *testing.T) {
	var a, ainv, r fe
	feFromSmall(&a, 42)
	feInv(&ainv, &a)
	montMul(&r, &a, &ainv)
	if feEqual(&r, &feOne) != 1 {
		t.Fatal("42 * 42^{-1} != 1")
	}
}

func TestFieldSqrt(t *testing.T) {
	// Compute 7^2 = 49, then sqrt(49) should be ±7.
	var seven, square, root, neg7 fe
	feFromSmall(&seven, 7)
	montSquare(&square, &seven)

	if !feSqrt(&root, &square) {
		t.Fatal("sqrt(49) failed")
	}

	// root should be 7 or P-7.
	feNeg(&neg7, &seven)
	if feEqual(&root, &seven) != 1 && feEqual(&root, &neg7) != 1 {
		t.Fatal("sqrt(49) != ±7")
	}
}

func TestFieldBytes27RoundTrip(t *testing.T) {
	var a, b fe
	feFromSmall(&a, 12345)
	var buf [27]byte
	feToBytes27(buf[:], &a)

	if !feFromBytes27(&b, buf[:]) {
		t.Fatal("feFromBytes27 rejected valid input")
	}
	if feEqual(&a, &b) != 1 {
		t.Fatal("round-trip failed")
	}
}

func TestScalarArithmetic(t *testing.T) {
	a := scalar{7, 0, 0, 0}
	b := scalar{13, 0, 0, 0}

	// 7 + 13 = 20.
	var sum scalar
	scAdd(&sum, &a, &b)
	if sum[0] != 20 || sum[1] != 0 || sum[2] != 0 || sum[3] != 0 {
		t.Fatalf("7 + 13 = %v, want 20", sum)
	}

	// 13 - 7 = 6.
	var diff scalar
	scSub(&diff, &b, &a)
	if diff[0] != 6 || diff[1] != 0 || diff[2] != 0 || diff[3] != 0 {
		t.Fatalf("13 - 7 = %v, want 6", diff)
	}

	// 7 * 13 = 91.
	var prod scalar
	scMul(&prod, &a, &b)
	if prod[0] != 91 || prod[1] != 0 || prod[2] != 0 || prod[3] != 0 {
		t.Fatalf("7 * 13 = %v, want 91", prod)
	}
}

func TestScalarBytes27RoundTrip(t *testing.T) {
	a := scalar{0xdeadbeef, 0xcafebabe, 0x12345678, 0}
	var buf [27]byte
	scToBytes27(buf[:], &a)

	var b scalar
	scFromBytes27(&b, buf[:])
	if a != b {
		t.Fatalf("round-trip: got %v, want %v", b, a)
	}
}

func TestGeneratorOrder(t *testing.T) {
	// SchnorrGen^Q should be the identity.
	_, q, _ := Params()

	var m4gen, result mat4
	tmToMat4(&m4gen, &schnorrGenTM)
	m4PowBig(&result, &m4gen, q)

	if !m4IsIdentity(&result) {
		t.Fatal("SchnorrGen^Q != I")
	}
}

func TestGeneratorNotIdentity(t *testing.T) {
	if tmIsIdentity(&schnorrGenTM) {
		t.Fatal("SchnorrGen is identity")
	}
}

func TestGeneratorDetOne(t *testing.T) {
	// det(SchnorrGen) should be 1.
	var m4gen mat4
	tmToMat4(&m4gen, &schnorrGenTM)
	var det fe
	m4Det(&det, &m4gen)
	if feEqual(&det, &feOne) != 1 {
		t.Fatal("det(SchnorrGen) != 1")
	}
}

func TestKeyGeneration(t *testing.T) {
	sk, pk, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	// Private key should be non-zero.
	if scIsZero(&sk.s) {
		t.Fatal("zero private key")
	}

	// Public key should not be identity.
	if tmIsIdentity(&pk.tm) {
		t.Fatal("identity public key")
	}

	// Public key should have det = 1.
	var m4pk mat4
	tmToMat4(&m4pk, &pk.tm)
	var det fe
	m4Det(&det, &m4pk)
	if feEqual(&det, &feOne) != 1 {
		t.Fatal("det(PK) != 1")
	}
}

func TestKeySerializationRoundTrip(t *testing.T) {
	sk, _, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	// Private key round-trip.
	skBytes := sk.Bytes()
	if len(skBytes) != 27 {
		t.Fatalf("sk bytes = %d, want 27", len(skBytes))
	}
	sk2, err := PrivateKeyFromBytes(skBytes)
	if err != nil {
		t.Fatal(err)
	}
	if sk.s != sk2.s {
		t.Logf("sk.s  = %016x %016x %016x %016x", sk.s[3], sk.s[2], sk.s[1], sk.s[0])
		t.Logf("sk2.s = %016x %016x %016x %016x", sk2.s[3], sk2.s[2], sk2.s[1], sk2.s[0])
		t.Logf("bytes = %x", skBytes)
		t.Fatal("private key round-trip failed")
	}
}

// mockGMid is a stand-in challenge hash for testing without importing pkg/crypto.
func mockGMid(data []byte) [27]byte {
	var result [27]byte
	// FNV-1a-like hash spread across 27 bytes.
	var state uint64 = 14695981039346656037
	for _, b := range data {
		state ^= uint64(b)
		state *= 1099511628211
	}
	for i := range result {
		state ^= uint64(i)
		state *= 1099511628211
		result[i] = byte(state >> 32)
	}
	return result
}

func TestSignVerifyRoundTrip(t *testing.T) {
	sk, pk, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	msg := []byte("test message for gnarl signature")

	sig, err := Sign(sk, msg, mockGMid)
	if err != nil {
		t.Fatal(err)
	}

	if !Verify(pk, msg, sig, mockGMid) {
		t.Fatal("valid signature rejected")
	}
}

func TestSignVerifyTamperedMessage(t *testing.T) {
	sk, pk, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	msg := []byte("original message")
	sig, err := Sign(sk, msg, mockGMid)
	if err != nil {
		t.Fatal(err)
	}

	tampered := []byte("tampered message")
	if Verify(pk, tampered, sig, mockGMid) {
		t.Fatal("tampered message accepted")
	}
}

func TestSignVerifyWrongKey(t *testing.T) {
	sk, _, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	_, pk2, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	msg := []byte("test message")
	sig, err := Sign(sk, msg, mockGMid)
	if err != nil {
		t.Fatal(err)
	}

	if Verify(pk2, msg, sig, mockGMid) {
		t.Fatal("wrong key accepted")
	}
}

func TestSignatureSerialization(t *testing.T) {
	sk, _, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	msg := []byte("serialize me")
	sig, err := Sign(sk, msg, mockGMid)
	if err != nil {
		t.Fatal(err)
	}

	sigBytes := sig.Bytes()
	if len(sigBytes) != 54 {
		t.Fatalf("signature bytes = %d, want 54", len(sigBytes))
	}

	sig2, err := SignatureFromBytes(sigBytes)
	if err != nil {
		t.Fatal(err)
	}

	if sig.e != sig2.e {
		t.Fatal("challenge mismatch after round-trip")
	}
	if sig.z != sig2.z {
		t.Fatal("response mismatch after round-trip")
	}
}

func TestPublicKeyYBytesLength(t *testing.T) {
	_, pk, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	yBytes := pk.YBytes()
	if len(yBytes) != 27 {
		t.Fatalf("y-only pubkey = %d bytes, want 27", len(yBytes))
	}
}

func TestTorusMatrixDetInvariant(t *testing.T) {
	// Generate a few keys and check det = 1 for all intermediate operations.
	for i := 0; i < 5; i++ {
		sk, _, err := GenerateKey()
		if err != nil {
			t.Fatal(err)
		}

		var pkTM tmat
		tmFixedExp(&pkTM, &sk.s)

		var m4pk mat4
		tmToMat4(&m4pk, &pkTM)
		var det fe
		m4Det(&det, &m4pk)
		if feEqual(&det, &feOne) != 1 {
			t.Fatalf("iteration %d: det != 1", i)
		}
	}
}

// --- Benchmarks ---

func BenchmarkMontMul(b *testing.B) {
	var a, c fe
	feFromSmall(&a, 42)
	feFromSmall(&c, 7)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		montMul(&a, &a, &c)
	}
}

func BenchmarkMontSquare(b *testing.B) {
	var a fe
	feFromSmall(&a, 42)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		montSquare(&a, &a)
	}
}

func BenchmarkTmMul(b *testing.B) {
	sk, _, err := GenerateKey()
	if err != nil {
		b.Fatal(err)
	}
	var pk tmat
	tmFixedExp(&pk, &sk.s)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tmMul(&pk, &pk, &schnorrGenTM)
	}
}

func BenchmarkTmSquare(b *testing.B) {
	sk, _, err := GenerateKey()
	if err != nil {
		b.Fatal(err)
	}
	var pk tmat
	tmFixedExp(&pk, &sk.s)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tmSquare(&pk, &pk)
	}
}

func BenchmarkFixedExp(b *testing.B) {
	var s scalar
	scRandom(&s)
	var r tmat
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tmFixedExp(&r, &s)
	}
}

func BenchmarkShamirExp(b *testing.B) {
	sk, pk, err := GenerateKey()
	if err != nil {
		b.Fatal(err)
	}
	var s2 scalar
	scRandom(&s2)
	var r tmat
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tmShamirExp(&r, &sk.s, &pk.tm, &s2)
	}
}

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

func BenchmarkVerify(b *testing.B) {
	sk, pk, err := GenerateKey()
	if err != nil {
		b.Fatal(err)
	}
	msg := []byte("benchmark message for verification")
	sig, err := Sign(sk, msg, mockGMid)
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !Verify(pk, msg, sig, mockGMid) {
			b.Fatal("verification failed")
		}
	}
}