package nizk

import (
	"fmt"
	"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)

	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

	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] {
						Z2 := Z.Div(Zs[i])
						X2 := Curve.Identity()

						for k := range n {
							if i != k {
								X2 = X2.Mul(Xs[k])
							}
						}
						X2 = X2.Exp(bits[i][idx].Stage.x)
						// BUG TODO
						{
							b1 := X2.String()
							b2 := Z2.String()
							if strings.HasPrefix(b2, b1[:len(b1)-4]) {
								fmt.Println("stage1, winner by string")
								winner = i
							}
						}
						if X2.Equal(Z2) {
							fmt.Println("stage1, winner by equal")
							winner = i
						}
						tb.Logf("stage 1: testing max_idx %d, i %d, bit %d:\n%v vs %v", max_idx, i, idx, X2, Z2)
					}
				}
				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 {
					Z2 := Z.Div(Zs[i])
					X2 := Curve.Identity()

					for k := range n {
						if i != k {
							X2 = X2.Mul(Xs[k])
						}
					}
					X2 = X2.Exp(bits[i][idx].Stage.x)

					// BUG TODO
					{
						b1 := X2.String()
						b2 := Z2.String()
						if strings.HasPrefix(b2, b1[:len(b1)-4]) {
							fmt.Println("STAGE2, winner by string")
							winner = i
						}
					}
					if X2.Equal(Z2) {
						fmt.Println("STAGE2, winner by equal")

						winner = i
					}
					tb.Logf("STAGE 2: testing max_idx %d, i %d, bit %d:\n%v vs %v", max_idx, i, idx, X2, Z2)
				}
			}
			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)
	}
	fmt.Printf("winner: %d, max: %d, \nvals: %08b\n", winner, max, vals)
}

func TestVickrey4on6bit(t *testing.T) { runVickrey(4, 6, t) }