package commit

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

// This is a construction of a proof of a statement of the form
//    [(C = g^(ab))   && (A = g^a) && (Β = g^b)]
// || [(C = g^(ab+1)) && (A = g^a) && (Β = g^b)]
//
// for given C, A and B

type Statement struct {
	a    *Scalar
	b    *Scalar
	plus bool
	*Commitment
}

type Commitment struct {
	C *Point
	A *Point
	B *Point
}

func NewStatement(a, b *Scalar, plus bool) *Statement {
	return &Statement{
		a:          a,
		b:          b,
		plus:       plus,
		Commitment: commitment(a, b, plus),
	}
}

func commitment(a, b *Scalar, plus bool) *Commitment {
	var C *Point
	c := a.Mul(b)

	if plus {
		C = G.Exp(c.Add(One))
	} else {
		C = G.Exp(c)
	}
	return &Commitment{
		C: C,
		A: G.Exp(a),
		B: G.Exp(b),
	}
}

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

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

func (s *Statement) Proof() *Proof {
	var e [2][2]*Point
	var r1, r2, w *Scalar
	r1 = Curve.RandomScalar()
	r2 = Curve.RandomScalar()
	w = Curve.RandomScalar()

	if s.plus {
		e[0][0] = G.Exp(r1)
		e[0][1] = s.B.Exp(r1).Mul(G.Exp(w))
		e[1][0] = G.Exp(r2)
		e[1][1] = s.B.Exp(r2)
	} else {
		e[0][0] = G.Exp(r1)
		e[0][1] = s.B.Exp(r1)
		e[1][0] = G.Exp(r2).Mul(s.A.Exp(w))
		e[1][1] = s.B.Exp(r2).Mul(s.C.Div(G).Exp(w))
	}

	ch := Challenge(G, s.C, s.A, s.B, e[0][0], e[0][1], e[1][0], e[1][1])
	pr := &Proof{}

	if s.plus {
		pr.Ch[0] = w
		pr.Ch[1] = ch.Sub(w)
		pr.R[0] = r1.Sub(s.a.Mul(pr.Ch[0]))
		pr.R[1] = r2.Sub(s.a.Mul(pr.Ch[1]))
	} else {
		pr.Ch[0] = ch.Sub(w)
		pr.Ch[1] = w
		pr.R[0] = r1.Sub(s.a.Mul(pr.Ch[0]))
		pr.R[1] = r2
	}

	return pr
}

func (c *Commitment) Verify(p *Proof) bool {
	var e [2][2]*Point
	e[0][0] = G.Exp(p.R[0]).Mul(c.A.Exp(p.Ch[0]))
	e[0][1] = c.B.Exp(p.R[0]).Mul(c.C.Exp(p.Ch[0]))
	e[1][0] = G.Exp(p.R[1]).Mul(c.A.Exp(p.Ch[1]))
	e[1][1] = c.B.Exp(p.R[1]).Mul(c.C.Div(G).Exp(p.Ch[1]))
	ch := Challenge(G, c.C, c.A, c.B, e[0][0], e[0][1], e[1][0], e[1][1])
	return p.Ch[0].Add(p.Ch[1]).Equal(ch)
}