package nizk

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

type Bit struct {
	id  Bytes
	set bool
	α   *Scalar
	β   *Scalar

	*Commitment
	prf *Proof

	*Stage
}

type Commitment struct {
	A *Point // g^α
	B *Point // g^β
	C *Point // g^(ab)g^(set)
}

// This is a construction of a proof of a statement of the form
//
//	[(C = g^(αβ))   && (A = g^α) && (Β = g^β)]
//	 || [(C = g^(αβ+1)) && (A = g^α) && (Β = g^β)]
//
// for given C, A and B
type Proof struct {
	A *schnorr.Proof // Proof for knowledge of α in A = G^α
	B *schnorr.Proof // Proof for knowledge of β in B = G^β
	C struct {       // Proof for knowledge of statement above
		Ch [2]*Scalar
		R  [2]*Scalar
	}
}

func NewBit(id Bytes, set bool) (*Bit, *Commitment, *Proof) {
	α, β := Curve.RandomScalar(), Curve.RandomScalar()
	return NewBitFromScalars(id, set, α, β)
}

func NewBitFromScalars(id Bytes, set bool, α, β *Scalar) (*Bit, *Commitment, *Proof) {
	b := &Bit{
		id:  id,
		set: set,
		α:   α,
		β:   β,
	}
	return b, b.commit(), b.proof()
}

func (b *Bit) IsSet() bool {
	return b.set
}

func (b *Bit) commit() *Commitment {
	if b.Commitment != nil {
		return b.Commitment
	}

	var C *Point
	c := b.α.Mul(b.β)

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

func (s *Bit) proof() *Proof {
	if s.prf != nil {
		return s.prf
	}

	var e [2][2]*Point
	var r1, r2, w *Scalar
	r1 = Curve.RandomScalar()
	r2 = Curve.RandomScalar()
	w = Curve.RandomScalar()
	c := s.commit()

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

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

	if s.set {
		pr.C.Ch[0] = w
		pr.C.Ch[1] = ch.Sub(w)
		pr.C.R[0] = r1.Sub(s.α.Mul(pr.C.Ch[0]))
		pr.C.R[1] = r2.Sub(s.α.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.α.Mul(pr.C.Ch[0]))
		pr.C.R[1] = r2
	}
	pr.A = (*schnorr.Statement)(s.α).Proof(s.id)
	pr.B = (*schnorr.Statement)(s.β).Proof(s.id)

	s.prf = pr
	return pr
}

func (c *Commitment) Verify(id Bytes, p *Proof) bool {
	var e [2][2]*Point

	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], id)
	return p.C.Ch[0].Add(p.C.Ch[1]).Equal(ch) &&
		(*schnorr.Commitment)(c.A).Verify(p.A, id) &&
		(*schnorr.Commitment)(c.B).Verify(p.B, id)
}