package nizk

import (
	"fmt"

	. "kesim.org/seal/common"
)

type Stage struct {
	x *Scalar
	r *Scalar

	*StageCommitment
	*StageReveal
}

type StageCommitment struct {
	R *Point
	X *Point
}

type StageReveal struct {
	Y *Point
	Z *Point
}

// Represents the proof 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 Stage1Proof struct {
	Ch  [2]*Scalar
	Rho [2][2]*Scalar
}

func (b *Bit) stage(x, r *Scalar) {
	b.Commit() // ensure non-null values for A, B, C
	b.Stage = &Stage{
		x: x,
		r: r,
	}
}

func (s *Stage) commit(Xs ...*Point) *StageCommitment {
	if s.StageCommitment != nil {
		return s.StageCommitment
	}

	s.StageCommitment = &StageCommitment{
		X: G.Exp(s.x),
		R: G.Exp(s.r),
	}
	return s.StageCommitment
}

func (b *Bit) StageCommit(Xs ...*Point) (s *StageCommitment) {
	x := Curve.RandomScalar()
	r := Curve.RandomScalar()
	return b.StageFromScalars(x, r, Xs...)
}

func (b *Bit) StageFromScalars(x, r *Scalar, Xs ...*Point) (c *StageCommitment) {
	b.stage(x, r)
	return b.Stage.commit(Xs...)
}

func (b *Bit) reveal(prev_true bool, Xs ...*Point) (r *StageReveal, e error) {
	s := b.Stage
	if s == nil {
		return nil, fmt.Errorf("stage not ready")
	}

	// TODO: Calculate Y based on the Xs and our own X_i
	// as Π_(i<k) X_k / Π_(i>k) X_k
	// For now:
	Y := G.Exp(Curve.RandomScalar())

	r = &StageReveal{Y: Y}

	if prev_true && b.IsSet() {
		r.Z = s.R.Exp(s.x)
	} else {
		r.Z = Y.Exp(s.x)
	}

	return r, e
}

func (b *Bit) RevealStage1(Xs ...*Point) (rev *StageReveal, pr *Stage1Proof, e error) {
	s := b.Stage
	if s == nil {
		return nil, nil, fmt.Errorf("stage not ready")
	}
	var ε [2][4]*Point
	var r1, r2, ρ1, ρ2, ω *Scalar
	for _, s := range []**Scalar{&r1, &r2, &ρ1, &ρ2, &ω} {
		*s = Curve.RandomScalar()
	}
	c := s.commit()

	rev, e = b.reveal(true, Xs...)
	if e != nil {
		return nil, nil, e
	}

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

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

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

	if b.IsSet() {
		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(b.α.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(b.α.Mul(pr.Ch[0]))
		pr.Rho[1][0] = ρ1
		pr.Rho[1][1] = ρ2
	}

	s.StageReveal = rev
	return rev, pr, e
}

func (c *Commitment) VerifyStage1(sc *StageCommitment, r *StageReveal, p *Stage1Proof) bool {
	var ε [2][4]*Point

	ε[0][0] = G.Exp(p.Rho[0][0]).Mul(sc.X.Exp(p.Ch[0]))
	ε[0][1] = G.Exp(p.Rho[0][1]).Mul(c.A.Exp(p.Ch[0]))
	ε[0][2] = r.Y.Exp(p.Rho[0][0]).Mul(r.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(sc.X.Exp(p.Ch[1]))
	ε[1][1] = G.Exp(p.Rho[1][1]).Mul(c.A.Exp(p.Ch[1]))
	ε[1][2] = sc.R.Exp(p.Rho[1][0]).Mul(r.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, sc.R, sc.X, r.Y, r.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)
}