package stage1

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

// Implements the proof and verification of statements of the following form:
//     [ Z=g^(xy) && X=g^x && Y=g^y && C=g^(αβ)   && A=g^α && B=g^β ]
//	|| [ Z=g^(xr) && X=g^x && R=g^r && C=g^(αβ+1) && A=g^α && B=g^β ]
// for given Z, X, Y, R, C, A and B

type Statement struct {
	x      *Scalar
	y      *Scalar
	r      *Scalar
	α      *Scalar
	β      *Scalar
	bitSet bool
	*Commitment
}

type Commitment struct {
	A *Point
	B *Point
	C *Point
	R *Point
	X *Point
	Y *Point
	Z *Point
}

func NewStatement(bitSet bool) *Statement {
	var x [5]*Scalar
	for i := range x {
		x[i] = Curve.RandomScalar()
	}
	return NewStatementFromScalars(bitSet, x[0], x[1], x[2], x[3], x[4])
}

func NewStatementFromScalars(bitSet bool, x, y, r, α, β *Scalar) *Statement {
	return &Statement{
		x:          x,
		y:          y,
		r:          r,
		α:          α,
		β:          β,
		bitSet:     bitSet,
		Commitment: commitment(x, y, r, α, β, bitSet),
	}
}

func commitment(x, y, r, α, β *Scalar, bitSet bool) *Commitment {
	var Z *Point
	φ := α.Mul(β)
	if bitSet {
		Z = G.Exp(x.Mul(r))
		φ = φ.Add(One)
	} else {
		Z = G.Exp(x.Mul(y))
	}

	return &Commitment{
		Z: Z,
		X: G.Exp(x),
		Y: G.Exp(y),
		R: G.Exp(r),
		A: G.Exp(α),
		B: G.Exp(β),
		C: G.Exp(φ),
	}
}

func (s *Statement) Commit() *Commitment {
	return s.Commitment
}

type Proof struct {
	Ch  [2]*Scalar
	Rho [2][2]*Scalar
}

func (s *Statement) Proof() *Proof {
	var ε [2][4]*Point
	var r1, r2, ρ1, ρ2, ω *Scalar
	for _, s := range []**Scalar{&r1, &r2, &ρ1, &ρ2, &ω} {
		*s = Curve.RandomScalar()
	}

	if s.bitSet {
		ε[0][0] = G.Exp(r1).Mul(s.X.Exp(ω))
		ε[0][1] = G.Exp(r2).Mul(s.A.Exp(ω))
		ε[0][2] = s.Y.Exp(r1).Mul(s.Z.Exp(ω))
		ε[0][3] = s.B.Exp(r2).Mul(s.C.Exp(ω))
		ε[1][0] = G.Exp(ρ1)
		ε[1][1] = G.Exp(ρ2)
		ε[1][2] = s.R.Exp(ρ1)
		ε[1][3] = s.B.Exp(ρ2)
	} else {
		ε[0][0] = G.Exp(r1)
		ε[0][1] = G.Exp(r2)
		ε[0][2] = s.Y.Exp(r1)
		ε[0][3] = s.B.Exp(r2)
		ε[1][0] = G.Exp(ρ1).Mul(s.X.Exp(ω))
		ε[1][1] = G.Exp(ρ2).Mul(s.A.Exp(ω))
		ε[1][2] = s.R.Exp(ρ1).Mul(s.Z.Exp(ω))
		ε[1][3] = s.B.Exp(ρ2).Mul(s.C.Div(G).Exp(ω))
	}

	p := []Bytes{G, s.A, s.B, s.C, s.R, s.X, s.Y, s.Z}
	for _, e := range ε[0] {
		p = append(p, e)
	}
	for _, e := range ε[1] {
		p = append(p, e)
	}

	ch := Challenge(p...)
	pr := &Proof{}

	if s.bitSet {
		pr.Ch[0] = ω
		pr.Ch[1] = ch.Sub(ω)
		pr.Rho[0][0] = r1
		pr.Rho[0][1] = r2
		pr.Rho[1][0] = ρ1.Sub(s.x.Mul(pr.Ch[1]))
		pr.Rho[1][1] = ρ2.Sub(s.α.Mul(pr.Ch[1]))
	} else {
		pr.Ch[0] = ch.Sub(ω)
		pr.Ch[1] = ω
		pr.Rho[0][0] = r1.Sub(s.x.Mul(pr.Ch[0]))
		pr.Rho[0][1] = r2.Sub(s.α.Mul(pr.Ch[0]))
		pr.Rho[1][0] = ρ1
		pr.Rho[1][1] = ρ2
	}

	return pr
}

func (c *Commitment) Verify(p *Proof) bool {
	var ε [2][4]*Point

	ε[0][0] = G.Exp(p.Rho[0][0]).Mul(c.X.Exp(p.Ch[0]))
	ε[0][1] = G.Exp(p.Rho[0][1]).Mul(c.A.Exp(p.Ch[0]))
	ε[0][2] = c.Y.Exp(p.Rho[0][0]).Mul(c.Z.Exp(p.Ch[0]))
	ε[0][3] = c.B.Exp(p.Rho[0][1]).Mul(c.C.Exp(p.Ch[0]))
	ε[1][0] = G.Exp(p.Rho[1][0]).Mul(c.X.Exp(p.Ch[1]))
	ε[1][1] = G.Exp(p.Rho[1][1]).Mul(c.A.Exp(p.Ch[1]))
	ε[1][2] = c.R.Exp(p.Rho[1][0]).Mul(c.Z.Exp(p.Ch[1]))
	ε[1][3] = c.B.Exp(p.Rho[1][1]).Mul(c.C.Div(G).Exp(p.Ch[1]))

	points := []Bytes{G, c.A, c.B, c.C, c.R, c.X, c.Y, c.Z}
	for _, e := range ε[0] {
		points = append(points, e)
	}
	for _, e := range ε[1] {
		points = append(points, e)
	}

	ch := Challenge(points...)

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