package nizk

import (
	"math/rand"
	"slices"
	"sync"
	"testing"

	. "kesim.org/seal/common"
)

func BenchmarkFromPaper(b *testing.B) {
	bitlength := 5
	vals := []int{
		0b01010,
		0b01001,
		0b00111,
	}
	ids := []Bytes{
		Curve.RandomScalar(),
		Curve.RandomScalar(),
		Curve.RandomScalar(),
	}

	for range b.N {
		var bits = [3][]*Bit{}
		for i, b := range vals {
			bits[i] = Int2Bits(ids[i], 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 {
				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)
					if !bits[i][idx].Commitment.VerifyStage1(c[i], r[i], p[i]) {
						b.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) {
					junction = idx
					instage1 = false

					for i := range bits {
						if !lost[i] && !bits[i][idx].IsSet() {
							lost[i] = true
						}
					}
					result |= 1 << (bitlength - 1 - idx)
				}
			} else {
				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)
					if !bits[i][idx].Commitment.VerifyStage2(bj[i].StageCommitment, c[i], bj[i].StageReveal, r[i], p[i]) {
						b.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) {
					junction = idx

					for i := range bits {
						if !lost[i] && !bits[i][idx].IsSet() {
							lost[i] = true
						}
					}
					result |= 1 << (bitlength - 1 - idx)
				}
			}
		}
		if result != vals[0] {
			b.Fatalf("wrong result: %05b, expected: %05b", result, vals[0])
		}
	}
}

func runSeal(n int, bitlength int, tb testing.TB) {
	var vals = make([]int, n)
	var ids = make([]Bytes, n)
	for i := range n {
		vals[i] = rand.Intn(1<<(bitlength-1) - 1)
		ids[i] = Curve.RandomScalar()
	}
	max := slices.Max(vals)

	var bits = make([][]*Bit, n)
	for i, b := range vals {
		bits[i] = Int2Bits(ids[i], b, bitlength)
	}

	var c = make([]*StageCommitment, n)
	var Xs = make([]*Point, n)
	var r = make([]*StageReveal, n)
	var p1 = make([]*Stage1Proof, n)
	var Zs = make([]*Point, n)
	var bj = make([]*Bit, n)
	var p2 = make([]*Stage2Proof, n)
	var lost = make([]bool, n)
	stage := stage1
	junction := -1
	result := 0

	for idx := range bitlength {
		for i := range n {
			c[i] = bits[i][idx].StageCommit()
			Xs[i] = c[i].X
			if stage == stage2 {
				bj[i] = bits[i][junction]

			}
		}
		var wg sync.WaitGroup
		wg.Add(n)

		for i := range n {
			go func() {
				defer wg.Done()

				if stage == stage1 {
					r[i], p1[i] = bits[i][idx].RevealStage1(Xs...)
					if !bits[i][idx].Commitment.VerifyStage1(c[i], r[i], p1[i]) {
						tb.Fatalf("bits[%d][%d] commitment failed to verify in stage1", i, idx)
					}
				} else {
					r[i], p2[i] = bits[i][idx].RevealStage2(lost[i], bj[i], Xs...)
					if !bits[i][idx].Commitment.VerifyStage2(bj[i].StageCommitment, c[i], bj[i].StageReveal, r[i], p2[i]) {
						tb.Fatalf("bits[%d][%d] commitment failed to verify in stage2, result so far: %05b", i, idx, result)
					}
				}
				Zs[i] = r[i].Z
			}()
		}
		wg.Wait()

		Z := Curve.Product(Zs...)
		if !Id.Equal(Z) {
			junction = idx
			stage = stage2
			for i := range bits {
				if !lost[i] && !bits[i][idx].IsSet() {
					lost[i] = true
				}
			}
			result |= 1 << (bitlength - 1 - idx)
		}
	}
	if result != max {
		tb.Fatalf("wrong result: %0[1]*[2]b, expected: %0[1]*[3]b", bitlength, result, max)
	}
}

func benchmarkMulti(n int, bitlength int, b *testing.B) {
	for range b.N {
		runSeal(n, bitlength, b)
	}
}

func BenchmarkRun3on5bit(b *testing.B)     { benchmarkMulti(3, 5, b) }
func BenchmarkRun10on8bit(b *testing.B)    { benchmarkMulti(10, 8, b) }
func BenchmarkRun10on16bit(b *testing.B)   { benchmarkMulti(10, 16, b) }
func BenchmarkRun10on24bit(b *testing.B)   { benchmarkMulti(10, 24, b) }
func BenchmarkRun100on8bit(b *testing.B)   { benchmarkMulti(100, 8, b) }
func BenchmarkRun100on16bit(b *testing.B)  { benchmarkMulti(100, 16, b) }
func BenchmarkRun100on24bit(b *testing.B)  { benchmarkMulti(100, 24, b) }
func BenchmarkRun500on8bit(b *testing.B)   { benchmarkMulti(500, 8, b) }
func BenchmarkRun500on16bit(b *testing.B)  { benchmarkMulti(500, 16, b) }
func BenchmarkRun500on24bit(b *testing.B)  { benchmarkMulti(500, 24, b) }
func BenchmarkRun1000on8bit(b *testing.B)  { benchmarkMulti(1000, 8, b) }
func BenchmarkRun1000on16bit(b *testing.B) { benchmarkMulti(1000, 16, b) }
func BenchmarkRun1000on24bit(b *testing.B) { benchmarkMulti(1000, 24, b) }