package commit

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

// 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
	Proof *Proof
}

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(id Bytes) *Commitment {
	s.Commitment.Proof = s.Proof(id)
	return s.Commitment
}

type Proof struct {
	Id Bytes
	A  *schnorr.Proof // Proof for knowledge of a in A = G^a
	B  *schnorr.Proof // Proof for knowledge of b in B = G^b
	C  struct {       // Proof for knowledge of statement above
		Ch [2]*Scalar
		R  [2]*Scalar
	}
}

func (s *Statement) Proof(id Bytes) *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], id)
	pr := &Proof{Id: id}

	if s.plus {
		pr.C.Ch[0] = w
		pr.C.Ch[1] = ch.Sub(w)
		pr.C.R[0] = r1.Sub(s.a.Mul(pr.C.Ch[0]))
		pr.C.R[1] = r2.Sub(s.a.Mul(pr.C.Ch[1]))
	} else {
		pr.C.Ch[0] = ch.Sub(w)
		pr.C.Ch[1] = w
		pr.C.R[0] = r1.Sub(s.a.Mul(pr.C.Ch[0]))
		pr.C.R[1] = r2
	}
	pr.A = (*schnorr.Statement)(s.a).Proof(id)
	pr.B = (*schnorr.Statement)(s.b).Proof(id)

	return pr
}

func (c *Commitment) Verify() bool {
	var e [2][2]*Point
	p := c.Proof

	e[0][0] = G.Exp(p.C.R[0]).Mul(c.A.Exp(p.C.Ch[0]))
	e[0][1] = c.B.Exp(p.C.R[0]).Mul(c.C.Exp(p.C.Ch[0]))
	e[1][0] = G.Exp(p.C.R[1]).Mul(c.A.Exp(p.C.Ch[1]))
	e[1][1] = c.B.Exp(p.C.R[1]).Mul(c.C.Div(G).Exp(p.C.Ch[1]))
	ch := Challenge(G, c.C, c.A, c.B, e[0][0], e[0][1], e[1][0], e[1][1], p.Id)
	return p.C.Ch[0].Add(p.C.Ch[1]).Equal(ch) &&
		(*schnorr.Commitment)(c.A).Verify(p.A, p.Id) &&
		(*schnorr.Commitment)(c.B).Verify(p.B, p.Id)
}