package nizk

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

	. "kesim.org/seal/common"
)

func runVickrey(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)
	max_idx := slices.Index(vals, max)
	if max_idx < len(vals)-1 && max == slices.Max(vals[max_idx+1:]) {
		max_idx = -1
	}

	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)
	instage1 := true
	junction := -1
	winner := -1
	result := 0

	isWinner := func(Z *Point, i int, idx int) bool {
		Z2 := Z.Div(Zs[i])
		xu := Curve.Identity()
		xl := Curve.Identity()

		for k := range n {
			if k < i {
				xu = xu.Mul(Xs[k])
			} else if k > i {
				xl = xl.Div(Xs[k])
			}
		}
		xu = xu.Exp(bits[i][idx].Stage.x)
		xl = xl.Exp(bits[i][idx].Stage.x)
		X2 := xu.Mul(xl)

		if X2.Equal(Z2) {
			tb.Logf("equal by value")
			winner = i
		}

		if winner < 0 {
			// BUG! TODO!
			s1 := X2.String()
			s2 := Z2.String()
			if strings.HasPrefix(s1, s2[:len(s2)-2]) {
				log.Printf("BUG! TODO! equal only by string")
				winner = i
			}

		}
		tb.Logf("testing max_idx %d, i %d, bit %d:\n%v vs %v", max_idx, i, idx, X2, Z2)
		return winner >= 0
	}

	for idx := range bitlength {
		for i := range n {
			c[i] = bits[i][idx].StageCommit()
			Xs[i] = c[i].X
		}

		if instage1 {
			var wg sync.WaitGroup
			wg.Add(n)

			for i := range n {
				go func() {
					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)
					}
					Zs[i] = r[i].Z
					wg.Done()
				}()
			}
			wg.Wait()

			Z := Curve.Product(Zs...)
			if !Id.Equal(Z) {
				junction = idx
				instage1 = false
				for i := range bits {
					if !lost[i] && !bits[i][idx].IsSet() {
						lost[i] = true
					}

					// Winner test
					if !lost[i] {
						isWinner(Z, i, idx)
					}
				}
				result |= 1 << (bitlength - 1 - idx)
			}
			continue
		}

		// Stage 2

		for i := range bits {
			bj[i] = bits[i][junction]
		}

		var wg sync.WaitGroup
		wg.Add(n)
		for i := range n {
			go func() {
				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.Done()
			}()
		}
		wg.Wait()

		Z := Curve.Product(Zs...)
		if !Id.Equal(Z) {
			junction = idx
			for i := range n {
				if !lost[i] && !bits[i][idx].IsSet() {
					lost[i] = true
				}

				// Winner test
				if !lost[i] && winner < 0 {
					isWinner(Z, i, idx)
				}
			}
			result |= 1 << (bitlength - 1 - idx)
		}
	}
	if result != max {
		tb.Fatalf("wrong result: %0[1]*[2]b, expected: %0[1]*[3]b", bitlength, result, max)
	}
	if max_idx != winner {
		tb.Fatalf("wrong winner, max_idx: %d vs winner: %d val %08b\nvals=%08b", max_idx, winner, max, vals)
	}
}

func TestVickrey100on24bit(t *testing.T) { runVickrey(100, 32, t) }