package nizk

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

func (b *Bit) CommitStage2(lost bool, prev *Stage) (s *Stage, c *StageCommitment, p *Stage2Proof) {
	x := Curve.RandomScalar()
	r := Curve.RandomScalar()
	return b.CommitStage2FromScalars(lost, prev, x, r)
}

func (b *Bit) CommitStage2FromScalars(lost bool, prev *Stage, x, r *Scalar) (s *Stage, c *StageCommitment, p *Stage2Proof) {
	s = b.stage(x, r)
	c = s.commit(lost)
	p = s.proof2(lost, prev)
	return
}

// 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 (s *Stage) proof2(lost bool, prev *Stage) *Stage2Proof {
	var (
		e1, e1_ [3]Bytes
		e2, e2_ [3]Bytes
		e3, e3_ [2]Bytes

		r1, r2 [3]*Scalar
		r3     [2]*Scalar
		w      [2]*Scalar
	)

	for _, scs := range [][]*Scalar{r1[:], r2[:], r3[:], w[:]} {
		for i := range scs {
			scs[i] = Curve.RandomScalar()
		}
	}

	c := s.commit(lost)
	bc := prev.bit.com
	pc := prev.com
	rvp := prev.rev
	rev := &StageReveal{
		Y: G.Exp(Curve.RandomScalar()), // TODO! BUG! THIS HAS TO BE Pj<i(X_j)/Pj>i(X_j)
	}
	if s.bit.IsSet() {
		rev.Z = c.R.Exp(s.x)
	} else {
		rev.Z = rev.Y.Exp(s.x)
	}

	if lost {
		e1[0] = G.Exp(r1[0]).Mul(c.X.Exp(w[0]))
		e1[1] = G.Exp(r1[1]).Mul(pc.X.Exp(w[0]))
		e1[2] = G.Exp(r1[2]).Mul(bc.A.Exp(w[0]))

		e1_[0] = c.R.Exp(r1[0]).Mul(rev.Z.Exp(w[0]))
		e1_[1] = pc.R.Exp(r1[1]).Mul(rvp.Z.Exp(w[0]))
		e1_[2] = bc.B.Exp(r1[2]).Mul(bc.C.Div(G).Exp(w[0]))

		e2[0] = G.Exp(r2[0]).Mul(c.X.Exp(w[1]))
		e2[1] = G.Exp(r2[1]).Mul(pc.X.Exp(w[1]))
		e2[2] = G.Exp(r2[2]).Mul(bc.A.Exp(w[1]))

		e2_[0] = rev.Y.Exp(r2[0]).Mul(rev.Z.Exp(w[1]))
		e2_[1] = pc.R.Exp(r2[1]).Mul(rvp.Z.Exp(w[1]))
		e2_[2] = bc.B.Exp(r2[2]).Mul(bc.C.Exp(w[1]))

		e3[0] = G.Exp(r3[0])
		e3[1] = G.Exp(r3[1])

		e3_[0] = rev.Y.Exp(r3[0])
		e3_[1] = rvp.Y.Exp(r3[1])
	} else {
		if s.bit.IsSet() {
			e1[0] = G.Exp(r1[0])
			e1[1] = G.Exp(r1[1])
			e1[2] = G.Exp(r1[2])

			e1_[0] = c.R.Exp(r1[0])
			e1_[1] = pc.R.Exp(r1[1])
			e1_[2] = bc.B.Exp(r1[2])

			e2[0] = G.Exp(r2[0]).Mul(c.X.Exp(w[0]))
			e2[1] = G.Exp(r2[1]).Mul(pc.X.Exp(w[0]))
			e2[2] = G.Exp(r2[2]).Mul(bc.A.Exp(w[0]))

			e2_[0] = rev.Y.Exp(r2[0]).Mul(rev.Z.Exp(w[0]))
			e2_[1] = pc.R.Exp(r2[1]).Mul(rvp.Z.Exp(w[0]))
			e2_[2] = bc.B.Exp(r2[2]).Mul(bc.C.Exp(w[0]))

			e3[0] = G.Exp(r3[0]).Mul(c.X.Exp(w[1]))
			e3[1] = G.Exp(r3[1]).Mul(pc.X.Exp(w[1]))

			e3_[0] = rev.Y.Exp(r3[0]).Mul(rev.Z.Exp(w[1]))
			e3_[1] = rvp.Y.Exp(r3[1]).Mul(rvp.Z.Exp(w[1]))
		} else {
			e1[0] = G.Exp(r1[0]).Mul(c.X.Exp(w[0]))
			e1[1] = G.Exp(r1[1]).Mul(pc.X.Exp(w[0]))
			e1[2] = G.Exp(r1[2]).Mul(bc.A.Exp(w[0]))

			e1_[0] = c.R.Exp(r1[0]).Mul(rev.Z.Exp(w[0]))
			e1_[1] = pc.R.Exp(r1[1]).Mul(rvp.Z.Exp(w[0]))
			e1_[2] = bc.B.Exp(r1[2]).Mul(bc.C.Div(G).Exp(w[0]))

			e2[0] = G.Exp(r2[0])
			e2[1] = G.Exp(r2[1])
			e2[2] = G.Exp(r2[2])

			e2_[0] = rev.Y.Exp(r2[0])
			e2_[1] = pc.R.Exp(r2[1])
			e2_[2] = bc.B.Exp(r2[2])

			e3[0] = G.Exp(r3[0]).Mul(c.X.Exp(w[1]))
			e3[1] = G.Exp(r3[1]).Mul(pc.X.Exp(w[1]))

			e3_[0] = rev.Y.Exp(r3[0]).Mul(rev.Z.Exp(w[1]))
			e3_[1] = rvp.Y.Exp(r3[1]).Mul(rvp.Z.Exp(w[1]))
		}
	}

	points := []Bytes{G, bc.A, bc.B, bc.C, c.R, c.X, rev.Y, rev.Z, pc.R, pc.X, rvp.Y, rvp.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...)
	pr := &Stage2Proof{}

	if lost {
		pr.Ch[0] = w[0]
		pr.Ch[1] = w[1]
		pr.Ch[2] = ch.Sub(w[0]).Sub(w[1])

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

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

		pr.R3[0] = r3[0].Sub(s.x.Mul(pr.Ch[2]))
		pr.R3[1] = r3[1].Sub(prev.x.Mul(pr.Ch[2]))
	} else {
		if s.bit.IsSet() {
			pr.Ch[0] = ch.Sub(w[0]).Sub(w[1])
			pr.Ch[1] = w[0]
			pr.Ch[2] = w[1]

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

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

			pr.R3[0] = r3[0]
			pr.R3[1] = r3[1]
		} else {
			pr.Ch[0] = w[0]
			pr.Ch[1] = ch.Sub(w[0]).Sub(w[1])
			pr.Ch[2] = w[1]

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

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

			pr.R3[0] = r3[0]
			pr.R3[1] = r3[1]
		}
	}

	s.prf2 = pr
	return pr
}

func (c *Commitment) VerifyStage2(pcom, ccom *StageCommitment, prev, crev *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(ccom.X.Exp(p.Ch[0]))
	e1[1] = G.Exp(p.R1[1]).Mul(pcom.X.Exp(p.Ch[0]))
	e1[2] = G.Exp(p.R1[2]).Mul(c.A.Exp(p.Ch[0]))

	e1_[0] = ccom.R.Exp(p.R1[0]).Mul(crev.Z.Exp(p.Ch[0]))
	e1_[1] = pcom.R.Exp(p.R1[1]).Mul(prev.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(ccom.X.Exp(p.Ch[1]))
	e2[1] = G.Exp(p.R2[1]).Mul(pcom.X.Exp(p.Ch[1]))
	e2[2] = G.Exp(p.R2[2]).Mul(c.A.Exp(p.Ch[1]))

	e2_[0] = crev.Y.Exp(p.R2[0]).Mul(crev.Z.Exp(p.Ch[1]))
	e2_[1] = pcom.R.Exp(p.R2[1]).Mul(prev.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(ccom.X.Exp(p.Ch[2]))
	e3[1] = G.Exp(p.R3[1]).Mul(pcom.X.Exp(p.Ch[2]))

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

	points := []Bytes{G, c.A, c.B, c.C, ccom.R, ccom.X, crev.Y, crev.Z, pcom.R, pcom.X, prev.Y, prev.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)
}