package nizk

import (
	"slices"
	"testing"

	. "kesim.org/seal/common"
)

func TestStage2Simple1(t *testing.T) {
	id := Curve.RandomScalar()

	for _, lost := range []bool{true, false} {
		b1 := NewBit(id, !lost)
		c1 := b1.StageCommit()
		r1, _ := b1.RevealStage1(c1.X)

		// Because the first index is a junction, any subsequent
		// combination of Bits must verify with 'lost' set to true
		// in the RevealStage2 calls.
		for _, s := range [][2]bool{
			{false, false},
			{true, false},
			{false, true},
			{true, true},
		} {
			b2 := NewBit(id, s[0])
			b3 := NewBit(id, s[1])
			b4 := NewBit(id, s[1]) // same as b3

			c2 := b2.StageCommit()
			c3 := b3.StageCommit()
			c4 := b4.StageCommit()

			r2, p2 := b2.RevealStage2(lost, b1, c1.X, c2.X, c3.X)
			if !b2.Commitment.VerifyStage2(c1, c2, r1, r2, p2) {
				t.Fatalf("failed to verify b2: %t b3: %t bc2/b1", s[0], s[1])
			}

			r3, p3 := b3.RevealStage2(lost, b1, c1.X, c2.X, c3.X)
			if !b3.Commitment.VerifyStage2(c1, c3, r1, r3, p3) {
				t.Fatalf("failed to verify b1: %t b3: %t bc3/b1", s[0], s[1])
			}

			r4, p4 := b4.RevealStage2(lost, b1, c1.X, c2.X, c3.X, c4.X)
			if !b4.Commitment.VerifyStage2(c1, c4, r1, r4, p4) {
				t.Fatalf("failed to verify b1: %t b4: %t bc4/b1", s[0], s[1])
			}
		}
	}
}

func int2bits(bid int, bitlength int) []*Bit {
	id := Curve.RandomScalar()

	bits := make([]*Bit, bitlength)
	for i := range bitlength {
		bits[i] = NewBit(id, (bid>>bitlength-i)&1 != 0)
	}
	return bits
}

func TestStage2Complex(t *testing.T) {
	bid1 := 0b0101
	bid2 := 0b0010
	t.Logf("testing bid1: %05b vs. bid2: %05b", bid1, bid2)

	bitlength := 4

	bits1 := int2bits(bid1, bitlength)
	bits2 := int2bits(bid2, bitlength)

	lost1 := false
	lost2 := false

	if len(bits1) != len(bits2) || len(bits1) != bitlength {
		t.Fatalf("oops")
	}

	instage1 := true
	junction := -1
	result := 0

	for c := 0; c < len(bits1); c++ {
		b1 := bits1[c]
		b2 := bits2[c]

		c1 := b1.StageCommit()
		c2 := b2.StageCommit()

		if instage1 {
			t.Logf("Testing bit b1[%d] = %t vs b2[%d] = %t", c, b1.IsSet(), c, b2.IsSet())

			r1, p1 := b1.RevealStage1(c1.X, c2.X)
			r2, p2 := b2.RevealStage1(c1.X, c2.X)

			if !b1.Commitment.VerifyStage1(c1, r1, p1) {
				t.Fatalf("b1 commitment failed to verify in stage1")
			}
			if !b2.Commitment.VerifyStage1(c2, r2, p2) {
				t.Fatalf("b2 commitment failed to verify in stage1")
			}

			Z := Curve.Product(r1.Z, r2.Z)
			if !Id.Equal(Z) {
				t.Logf("Aha! Z[%d] != Id, switch to stage2", c)
				junction = c
				instage1 = false

				if !lost1 && !b1.IsSet() {
					t.Logf("setting lost1 to true")
					lost1 = true
				}

				if !lost2 && !b2.IsSet() {
					t.Logf("setting lost2 to true")
					lost2 = true
				}
				result |= 1 << (bitlength - 1 - c)
			} else {
				t.Logf("Z[%d] == Id, staying in stage1", c)
			}
		} else {
			t.Logf("Testing bit b1[%d]∧lost1 = %t vs b2[%d]∧lost2 = %t", c, b1.IsSet() && !lost1, c, b2.IsSet() && !lost2)

			bj1 := bits1[junction]
			bj2 := bits2[junction]

			r1, p1 := b1.RevealStage2(lost1, bj1, c1.X, c2.X)
			r2, p2 := b2.RevealStage2(lost2, bj2, c1.X, c2.X)

			if !b1.Commitment.VerifyStage2(bj1.StageCommitment, c1, bj1.StageReveal, r1, p1) {
				t.Fatalf("b1 commitment failed to verify in stage1")
			}
			if !b2.Commitment.VerifyStage2(bj2.StageCommitment, c2, bj2.StageReveal, r2, p2) {
				t.Fatalf("b2 commitment failed to verify in stage1")
			}

			Z := Curve.Product(r1.Z, r2.Z)
			if !Id.Equal(Z) {
				t.Logf("Aha! Z[%d] != Id, new junction!", c)
				junction = c

				if !lost1 && !b1.IsSet() {
					t.Logf("setting lost1 to true")
					lost1 = true
				}

				if !lost2 && !b2.IsSet() {
					t.Logf("setting lost2 to true")
					lost2 = true
				}
				result |= 1 << (bitlength - 1 - c)
			}
		}
	}
	if result != bid1 {
		t.Fatalf("wrong result: %05b", result)
	}
}

func TestFromPaper(t *testing.T) {
	bitlength := 5
	vals := []int{
		0b01001,
		0b01010,
		0b00111,
	}

	t.Logf("testing bits: %05b, %05b, %05b", vals[0], vals[1], vals[2])

	var bits = [3][]*Bit{}
	for i, b := range vals {
		bits[i] = int2bits(b, bitlength)
	}

	var lost = [3]bool{}
	instage1 := true
	junction := -1
	result := 0

	for idx := 0; idx < bitlength; idx++ {
		var c = [3]*StageCommitment{}
		var r = [3]*StageReveal{}

		for i, b := range bits {
			c[i] = b[idx].StageCommit()
		}

		if instage1 {
			t.Logf("bit[%d] b1 = %t vs b2 = %t vs b3 = %t", idx,
				bits[0][idx].IsSet(), bits[1][idx].IsSet(), bits[2][idx].IsSet())

			var p = [3]*Stage1Proof{}
			for i := range bits {
				r[i], p[i] = bits[i][idx].RevealStage1(c[0].X, c[1].X, c[2].X)
				t.Logf("bits[%d][%d] has sent %t", i, idx, bits[i][idx].Sent)
				if !bits[i][idx].Commitment.VerifyStage1(c[i], r[i], p[i]) {
					t.Fatalf("bits[%d][%d] commitment failed to verify in stage1", i, idx)
				}
			}

			Z := Curve.Product(r[0].Z, r[1].Z, r[2].Z)
			if !Id.Equal(Z) {
				t.Logf("Aha! Z[%d] != Id, switch to stage2", idx)
				junction = idx
				instage1 = false

				for i := range bits {
					if !lost[i] && !bits[i][idx].IsSet() {
						lost[i] = true
						t.Logf("bit %d, set lost[%d] to true, so far: %v", idx, i, lost)
					}
				}
				result |= 1 << (bitlength - 1 - idx)
			} else {
				t.Logf("Z[%d] == Id, staying in stage1", idx)
			}
		} else {
			t.Logf("\nTesing bit[%d]:\n"+
				"set  0: %t\t1: %t\t2: %t\n"+
				"lost 0: %t\t1  %t\t2: %t\n",
				idx,
				bits[0][idx].IsSet(), bits[1][idx].IsSet(), bits[2][idx].IsSet(),
				lost[0], lost[1], lost[2])

			var bj = [3]*Bit{}
			for i := range bits {
				bj[i] = bits[i][junction]
			}

			var p = [3]*Stage2Proof{}
			for i := range bits {
				r[i], p[i] = bits[i][idx].RevealStage2(lost[i], bj[i], c[0].X, c[1].X, c[2].X)
				t.Logf("bits[%d][%d] has sent %t", i, idx, bits[i][idx].Sent)
				if !bits[i][idx].Commitment.VerifyStage2(bj[i].StageCommitment, c[i], bj[i].StageReveal, r[i], p[i]) {
					t.Fatalf("bits[%d][%d] commitment failed to verify in stage2, result so far: %05b", i, idx, result)
				}
			}

			Z := Curve.Product(r[0].Z, r[1].Z, r[2].Z)
			if !Id.Equal(Z) {
				t.Logf("Aha! Z[%d] != Id, new junction!", idx)
				junction = idx

				for i := range bits {
					if !lost[i] && !bits[i][idx].IsSet() {
						lost[i] = true
						t.Logf("bits[%d][%d], set lost[%d] to true, so far: %v", i, idx, i, lost)
					}
				}
				result |= 1 << (bitlength - 1 - idx)
			}
		}
	}
	max := slices.Max(vals)
	if result != max {
		t.Fatalf("wrong result: %05b, expected: %05b", result, max)
	}

}

/*
func TestRun3on5bit(t *testing.T)    { runSeal(3, 5, t) }
func TestRun10on16bit(t *testing.T)  { runSeal(10, 16, t) }
func TestRun100on16bit(t *testing.T) { runSeal(100, 16, t) }
*/