package nizk

import (
	. "kesim.org/seal/common"
)

// Represents the proof of a statement of the following form:
//
//	( Z=g^(x*y) && X=g^x && Y=g^y && Z_=g^(x_*y_) && X_=g^x_ && Y_=g^y_ )	// case "none"
//
// || ( Z=g^(x*y) && X=g^x && Y=g^y && Z_=g^(x_*r_) && X_=g^x_ && R_=g^r_ && C=g^(a*b)   && A=g^a && B=g^b ) // case "unset"
// || ( Z=g^(x*r) && X=g^x && R=g^r && Z_=g^(x_*r_) && X_=g^x_ && R_=g^r_ && C=g^(a*b+1) && A=g^a && B=g^b ) // case "set"
//
// for given A, B, C, R, X, Y, Z, R_, X_, Y_, Z_ on the curve
type Stage2Proof struct {
	Ch [3]*Scalar
	R1 [3]*Scalar
	R2 [3]*Scalar
	R3 [2]*Scalar
}

func (b *Bit) RevealStage2(lost bool, prev *Bit, Xs ...*Point) (rv2 *StageReveal, pr *Stage2Proof) {
	if b.Stage == nil {
		b.StageCommit()
	}
	s := b.Stage

	var (
		ε1, ε1_ [3]Bytes
		ε2, ε2_ [3]Bytes
		ε3, ε3_ [2]Bytes

		ρ1, ρ2 [3]*Scalar
		ρ3     [2]*Scalar
		ω      [2]*Scalar
	)

	for _, s := range [][]*Scalar{ρ1[:], ρ2[:], ρ3[:], ω[:]} {
		for i := range s {
			s[i] = Curve.RandomScalar()
		}
	}

	c1 := prev.StageCommitment
	c2 := s.StageCommitment
	rv1 := prev.StageReveal
	rv2 = b.reveal(prev.IsSet(), Xs...)

	if lost {
		ε1[0] = G.Exp(ρ1[0]).Mul(c2.X.Exp(ω[0]))
		ε1[1] = G.Exp(ρ1[1]).Mul(c1.X.Exp(ω[0]))
		ε1[2] = G.Exp(ρ1[2]).Mul(b.A.Exp(ω[0]))

		ε1_[0] = c2.R.Exp(ρ1[0]).Mul(rv2.Z.Exp(ω[0]))
		ε1_[1] = c1.R.Exp(ρ1[1]).Mul(rv1.Z.Exp(ω[0]))
		ε1_[2] = b.B.Exp(ρ1[2]).Mul(b.C.Div(G).Exp(ω[0]))

		ε2[0] = G.Exp(ρ2[0]).Mul(c2.X.Exp(ω[1]))
		ε2[1] = G.Exp(ρ2[1]).Mul(c1.X.Exp(ω[1]))
		ε2[2] = G.Exp(ρ2[2]).Mul(b.A.Exp(ω[1]))

		ε2_[0] = rv2.Y.Exp(ρ2[0]).Mul(rv2.Z.Exp(ω[1]))
		ε2_[1] = c1.R.Exp(ρ2[1]).Mul(rv1.Z.Exp(ω[1]))
		ε2_[2] = b.B.Exp(ρ2[2]).Mul(b.C.Exp(ω[1]))

		ε3[0] = G.Exp(ρ3[0])
		ε3[1] = G.Exp(ρ3[1])

		ε3_[0] = rv2.Y.Exp(ρ3[0])
		ε3_[1] = rv1.Y.Exp(ρ3[1])
	} else {
		if b.IsSet() {
			ε1[0] = G.Exp(ρ1[0])
			ε1[1] = G.Exp(ρ1[1])
			ε1[2] = G.Exp(ρ1[2])

			ε1_[0] = c2.R.Exp(ρ1[0])
			ε1_[1] = c1.R.Exp(ρ1[1])
			ε1_[2] = b.B.Exp(ρ1[2])

			ε2[0] = G.Exp(ρ2[0]).Mul(c2.X.Exp(ω[0]))
			ε2[1] = G.Exp(ρ2[1]).Mul(c1.X.Exp(ω[0]))
			ε2[2] = G.Exp(ρ2[2]).Mul(b.A.Exp(ω[0]))

			ε2_[0] = rv2.Y.Exp(ρ2[0]).Mul(rv2.Z.Exp(ω[0]))
			ε2_[1] = c1.R.Exp(ρ2[1]).Mul(rv1.Z.Exp(ω[0]))
			ε2_[2] = b.B.Exp(ρ2[2]).Mul(b.C.Exp(ω[0]))

			ε3[0] = G.Exp(ρ3[0]).Mul(c2.X.Exp(ω[1]))
			ε3[1] = G.Exp(ρ3[1]).Mul(c1.X.Exp(ω[1]))

			ε3_[0] = rv2.Y.Exp(ρ3[0]).Mul(rv2.Z.Exp(ω[1]))
			ε3_[1] = rv1.Y.Exp(ρ3[1]).Mul(rv1.Z.Exp(ω[1]))
		} else {
			ε1[0] = G.Exp(ρ1[0]).Mul(c2.X.Exp(ω[0]))
			ε1[1] = G.Exp(ρ1[1]).Mul(c1.X.Exp(ω[0]))
			ε1[2] = G.Exp(ρ1[2]).Mul(b.A.Exp(ω[0]))

			ε1_[0] = c2.R.Exp(ρ1[0]).Mul(rv2.Z.Exp(ω[0]))
			ε1_[1] = c1.R.Exp(ρ1[1]).Mul(rv1.Z.Exp(ω[0]))
			ε1_[2] = b.B.Exp(ρ1[2]).Mul(b.C.Div(G).Exp(ω[0]))

			ε2[0] = G.Exp(ρ2[0])
			ε2[1] = G.Exp(ρ2[1])
			ε2[2] = G.Exp(ρ2[2])

			ε2_[0] = rv2.Y.Exp(ρ2[0])
			ε2_[1] = c1.R.Exp(ρ2[1])
			ε2_[2] = b.B.Exp(ρ2[2])

			ε3[0] = G.Exp(ρ3[0]).Mul(c2.X.Exp(ω[1]))
			ε3[1] = G.Exp(ρ3[1]).Mul(c1.X.Exp(ω[1]))

			ε3_[0] = rv2.Y.Exp(ρ3[0]).Mul(rv2.Z.Exp(ω[1]))
			ε3_[1] = rv1.Y.Exp(ρ3[1]).Mul(rv1.Z.Exp(ω[1]))
		}
	}

	points := []Bytes{G, b.A, b.B, b.C, c2.R, c2.X, rv2.Y, rv2.Z, c1.R, c1.X, rv1.Y, rv1.Z}
	points = append(points, ε1[:]...)
	points = append(points, ε2[:]...)
	points = append(points, ε3[:]...)
	points = append(points, ε1_[:]...)
	points = append(points, ε2_[:]...)
	points = append(points, ε3_[:]...)

	ch := Challenge(points...)
	pr = &Stage2Proof{}

	if !prev.IsSet() {
		pr.Ch[0] = ω[0]
		pr.Ch[1] = ω[1]
		pr.Ch[2] = ch.Sub(ω[0]).Sub(ω[1])

		pr.R1[0] = ρ1[0]
		pr.R1[1] = ρ1[1]
		pr.R1[2] = ρ1[2]

		pr.R2[0] = ρ2[0]
		pr.R2[1] = ρ2[1]
		pr.R2[2] = ρ2[2]

		pr.R3[0] = ρ3[0].Sub(s.x.Mul(pr.Ch[2]))
		pr.R3[1] = ρ3[1].Sub(prev.x.Mul(pr.Ch[2]))
	} else {
		if b.IsSet() {
			pr.Ch[0] = ch.Sub(ω[0]).Sub(ω[1])
			pr.Ch[1] = ω[0]
			pr.Ch[2] = ω[1]

			pr.R1[0] = ρ1[0].Sub(s.x.Mul(pr.Ch[0]))
			pr.R1[1] = ρ1[1].Sub(prev.x.Mul(pr.Ch[0]))
			pr.R1[2] = ρ1[2].Sub(b.α.Mul(pr.Ch[0]))

			pr.R2[0] = ρ2[0]
			pr.R2[1] = ρ2[1]
			pr.R2[2] = ρ2[2]

			pr.R3[0] = ρ3[0]
			pr.R3[1] = ρ3[1]
		} else {
			pr.Ch[0] = ω[0]
			pr.Ch[1] = ch.Sub(ω[0]).Sub(ω[1])
			pr.Ch[2] = ω[1]

			pr.R1[0] = ρ1[0]
			pr.R1[1] = ρ1[1]
			pr.R1[2] = ρ1[2]

			pr.R2[0] = ρ2[0].Sub(s.x.Mul(pr.Ch[1]))
			pr.R2[1] = ρ2[1].Sub(prev.x.Mul(pr.Ch[1]))
			pr.R2[2] = ρ2[2].Sub(b.α.Mul(pr.Ch[1]))

			pr.R3[0] = ρ3[0]
			pr.R3[1] = ρ3[1]
		}
	}

	return rv2, pr
}

func (c *Commitment) VerifyStage2(c1, c2 *StageCommitment, r1, r2 *StageReveal, p *Stage2Proof) bool {
	var (
		e1, e1_ [3]Bytes
		e2, e2_ [3]Bytes
		e3, e3_ [2]Bytes
	)
	e1[0] = G.Exp(p.R1[0]).Mul(c2.X.Exp(p.Ch[0]))
	e1[1] = G.Exp(p.R1[1]).Mul(c1.X.Exp(p.Ch[0]))
	e1[2] = G.Exp(p.R1[2]).Mul(c.A.Exp(p.Ch[0]))

	e1_[0] = c2.R.Exp(p.R1[0]).Mul(r2.Z.Exp(p.Ch[0]))
	e1_[1] = c1.R.Exp(p.R1[1]).Mul(r1.Z.Exp(p.Ch[0]))
	e1_[2] = c.B.Exp(p.R1[2]).Mul(c.C.Div(G).Exp(p.Ch[0]))

	e2[0] = G.Exp(p.R2[0]).Mul(c2.X.Exp(p.Ch[1]))
	e2[1] = G.Exp(p.R2[1]).Mul(c1.X.Exp(p.Ch[1]))
	e2[2] = G.Exp(p.R2[2]).Mul(c.A.Exp(p.Ch[1]))

	e2_[0] = r2.Y.Exp(p.R2[0]).Mul(r2.Z.Exp(p.Ch[1]))
	e2_[1] = c1.R.Exp(p.R2[1]).Mul(r1.Z.Exp(p.Ch[1]))
	e2_[2] = c.B.Exp(p.R2[2]).Mul(c.C.Exp(p.Ch[1]))

	e3[0] = G.Exp(p.R3[0]).Mul(c2.X.Exp(p.Ch[2]))
	e3[1] = G.Exp(p.R3[1]).Mul(c1.X.Exp(p.Ch[2]))

	e3_[0] = r2.Y.Exp(p.R3[0]).Mul(r2.Z.Exp(p.Ch[2]))
	e3_[1] = r1.Y.Exp(p.R3[1]).Mul(r1.Z.Exp(p.Ch[2]))

	points := []Bytes{G, c.A, c.B, c.C, c2.R, c2.X, r2.Y, r2.Z, c1.R, c1.X, r1.Y, r1.Z}
	points = append(points, e1[:]...)
	points = append(points, e2[:]...)
	points = append(points, e3[:]...)
	points = append(points, e1_[:]...)
	points = append(points, e2_[:]...)
	points = append(points, e3_[:]...)

	ch := Challenge(points...)

	return p.Ch[0].Add(p.Ch[1]).Add(p.Ch[2]).Equal(ch)
}