package nizk

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

// Implements the proof and verification 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 Type int

const (
	None Type = iota
	Unset
	Set
)

type Stage2 struct {
	typ Type
	a   *Scalar
	b   *Scalar
	r   *Scalar
	x   *Scalar
	y   *Scalar
	r_  *Scalar
	x_  *Scalar
	y_  *Scalar
}

type Stage2Commitment struct {
	A  *Point
	B  *Point
	C  *Point
	R  *Point
	X  *Point
	Y  *Point
	Z  *Point
	R_ *Point
	X_ *Point
	Y_ *Point
	Z_ *Point
}

func NewStage2(typ Type) *Stage2 {
	var s [8]*Scalar
	for i := range s {
		s[i] = Curve.RandomScalar()
	}

	return NewStage2FromScalars(typ, s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7])
}

func NewStage2FromScalars(typ Type, a, b, r, x, y, r_, x_, y_ *Scalar) *Stage2 {
	if typ > Set || typ < None {
		panic("unknown type")
	}

	return &Stage2{
		typ: typ,
		a:   a,
		b:   b,
		r:   r,
		x:   x,
		y:   y,
		r_:  r_,
		x_:  x_,
		y_:  y_,
	}
}

func (s *Stage2) commitment() *Stage2Commitment {
	var Z, Z_ *Point
	c := s.a.Mul(s.b)

	switch s.typ {
	case None:
		Z = G.Exp(s.x.Mul(s.y))
		Z_ = G.Exp(s.x_.Mul(s.y_))
	case Set:
		Z = G.Exp(s.x.Mul(s.r))
		Z_ = G.Exp(s.x_.Mul(s.r_))
		c = c.Add(One)
	case Unset:
		Z = G.Exp(s.x.Mul(s.y))
		Z_ = G.Exp(s.x_.Mul(s.r_))
	default:
		panic("not possible")
	}

	return &Stage2Commitment{
		A:  G.Exp(s.a),
		B:  G.Exp(s.b),
		C:  G.Exp(c),
		R:  G.Exp(s.r),
		X:  G.Exp(s.x),
		Y:  G.Exp(s.y),
		Z:  Z,
		R_: G.Exp(s.r_),
		X_: G.Exp(s.x_),
		Y_: G.Exp(s.y_),
		Z_: Z_,
	}
}

func (s *Stage2) Commit() (*Stage2Commitment, *Stage2Proof) {
	c := s.commitment()
	return c, s.proof(c)
}

type Stage2Proof struct {
	Ch [3]*Scalar
	R1 [3]*Scalar
	R2 [3]*Scalar
	R3 [2]*Scalar
}

func (s *Stage2) proof(c *Stage2Commitment) *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()
		}
	}

	switch s.typ {
	case None:
		e1[0] = G.Exp(r1[0]).Mul(c.X.Exp(w[0]))
		e1[1] = G.Exp(r1[1]).Mul(c.X_.Exp(w[0]))
		e1[2] = G.Exp(r1[2]).Mul(c.A.Exp(w[0]))

		e1_[0] = c.R.Exp(r1[0]).Mul(c.Z.Exp(w[0]))
		e1_[1] = c.R_.Exp(r1[1]).Mul(c.Z_.Exp(w[0]))
		e1_[2] = c.B.Exp(r1[2]).Mul(c.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(c.X_.Exp(w[1]))
		e2[2] = G.Exp(r2[2]).Mul(c.A.Exp(w[1]))

		e2_[0] = c.Y.Exp(r2[0]).Mul(c.Z.Exp(w[1]))
		e2_[1] = c.R_.Exp(r2[1]).Mul(c.Z_.Exp(w[1]))
		e2_[2] = c.B.Exp(r2[2]).Mul(c.C.Exp(w[1]))

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

		e3_[0] = c.Y.Exp(r3[0])
		e3_[1] = c.Y_.Exp(r3[1])

	case Unset:
		e1[0] = G.Exp(r1[0]).Mul(c.X.Exp(w[0]))
		e1[1] = G.Exp(r1[1]).Mul(c.X_.Exp(w[0]))
		e1[2] = G.Exp(r1[2]).Mul(c.A.Exp(w[0]))

		e1_[0] = c.R.Exp(r1[0]).Mul(c.Z.Exp(w[0]))
		e1_[1] = c.R_.Exp(r1[1]).Mul(c.Z_.Exp(w[0]))
		e1_[2] = c.B.Exp(r1[2]).Mul(c.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] = c.Y.Exp(r2[0])
		e2_[1] = c.R_.Exp(r2[1])
		e2_[2] = c.B.Exp(r2[2])

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

		e3_[0] = c.Y.Exp(r3[0]).Mul(c.Z.Exp(w[1]))
		e3_[1] = c.Y_.Exp(r3[1]).Mul(c.Z_.Exp(w[1]))

	case Set:
		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] = c.R_.Exp(r1[1])
		e1_[2] = c.B.Exp(r1[2])

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

		e2_[0] = c.Y.Exp(r2[0]).Mul(c.Z.Exp(w[0]))
		e2_[1] = c.R_.Exp(r2[1]).Mul(c.Z_.Exp(w[0]))
		e2_[2] = c.B.Exp(r2[2]).Mul(c.C.Exp(w[0]))

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

		e3_[0] = c.Y.Exp(r3[0]).Mul(c.Z.Exp(w[1]))
		e3_[1] = c.Y_.Exp(r3[1]).Mul(c.Z_.Exp(w[1]))

	default:
		panic("not possible")
	}

	points := []Bytes{G, c.A, c.B, c.C, c.R, c.X, c.Y, c.Z, c.R_, c.X_, c.Y_, c.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{}

	switch s.typ {
	case None:
		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(s.x_.Mul(pr.Ch[2]))

	case Unset:
		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(s.x_.Mul(pr.Ch[1]))
		pr.R2[2] = r2[2].Sub(s.a.Mul(pr.Ch[1]))

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

	case Set:
		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(s.x_.Mul(pr.Ch[0]))
		pr.R1[2] = r1[2].Sub(s.a.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]

	default:
		panic("unreachable")
	}

	return pr
}

func (c *Stage2Commitment) Verify(p *Stage2Proof) bool {
	var (
		e1, e1_ [3]Bytes
		e2, e2_ [3]Bytes
		e3, e3_ [2]Bytes
	)
	e1[0] = G.Exp(p.R1[0]).Mul(c.X.Exp(p.Ch[0]))
	e1[1] = G.Exp(p.R1[1]).Mul(c.X_.Exp(p.Ch[0]))
	e1[2] = G.Exp(p.R1[2]).Mul(c.A.Exp(p.Ch[0]))

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

	e2_[0] = c.Y.Exp(p.R2[0]).Mul(c.Z.Exp(p.Ch[1]))
	e2_[1] = c.R_.Exp(p.R2[1]).Mul(c.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(c.X.Exp(p.Ch[2]))
	e3[1] = G.Exp(p.R3[1]).Mul(c.X_.Exp(p.Ch[2]))

	e3_[0] = c.Y.Exp(p.R3[0]).Mul(c.Z.Exp(p.Ch[2]))
	e3_[1] = c.Y_.Exp(p.R3[1]).Mul(c.Z_.Exp(p.Ch[2]))

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