package vote

import (
	"crypto/rand"
	"crypto/sha512"
	"encoding/base32"
	"encoding/binary"
	"fmt"
	"io"

	curve "filippo.io/edwards25519"
)

var b32 = base32.StdEncoding.WithPadding(base32.NoPadding)

// A Schnorr signature to prove knowledge of x for given g^x and i.

type Proof struct {
	PointV  curve.Point
	ScalarR curve.Scalar
	I       uint64
}

type Commitment struct {
	Index  uint64
	PubX   curve.Point
	PubR   curve.Point
	ProofX Proof
	ProofR Proof
}

type Vote struct {
	bit bool

	privX curve.Scalar
	privR curve.Scalar

	Commitment
}

func newPriv(s *curve.Scalar, random io.Reader) error {
	var buf [64]byte
	if random == nil {
		random = rand.Reader
	}
	random.Read(buf[:])
	_, e := s.SetUniformBytes(buf[:])
	return e
}

func setPub(p *curve.Scalar, P *curve.Point) *curve.Point {
	return P.ScalarBaseMult(p)
}

// Generates the proof, aka Schnorr signature, for given priv and i.
// Choosing a scalar v randomly, the signature consists of (V, r) with
//
//	V := g^v, with randomly chosen v
//	 r := (v - x*h), with h := H(g, g^v, g^x, i), where i is given by the context.
//
// Verification of the signature is by comparing V =?= g^r * g^(x*h)
func genProof(pr *Proof, x *curve.Scalar, i uint64) error {
	pr.I = i
	var v = new(curve.Scalar)
	e := newPriv(v, nil)
	if e != nil {
		return e
	}
	setPub(v, &pr.PointV)
	gx := new(curve.Point)
	setPub(x, gx)

	// Calculate h := H(g, g^v, g^x, i)
	h, e := hash(&pr.PointV, gx, i)

	if e != nil {
		return e
	}

	// Calculate r := v - x*h
	xh := new(curve.Scalar).Multiply(x, h)
	(&pr.ScalarR).Subtract(v, xh)

	return nil
}

// Calculate  h := H(g, g^v, g^x, i)
func hash(gv, gx *curve.Point, i uint64) (*curve.Scalar, error) {
	h512 := sha512.New()
	h512.Write(curve.NewGeneratorPoint().Bytes())
	h512.Write(gv.Bytes())
	h512.Write(gx.Bytes())
	_ = binary.Write(h512, binary.BigEndian, i)

	hb := h512.Sum(nil)
	return new(curve.Scalar).SetUniformBytes(hb)
}

func (v *Vote) genProofs() (e error) {
	e = genProof(&v.ProofX, &v.privX, v.Index)
	if e != nil {
		return e
	}
	return genProof(&v.ProofR, &v.privR, v.Index)
}

// Verifies that g^v == g^r*g^(x*h)
func verifyProof(V *curve.Point, r, x *curve.Scalar, i uint64) (ok bool) {
	// Calculate h = H(g, g^v, g^x, i)
	gx := new(curve.Point)
	setPub(x, gx)
	h, e := hash(V, gx, i)
	if e != nil {
		return false
	}

	// Calculate g^(x*h)
	xh := new(curve.Scalar).Multiply(x, h)
	gxh := new(curve.Point)
	setPub(xh, gxh)

	// Calculate g^r
	gr := new(curve.Point)
	setPub(r, gr)

	// Calculate g^r*g^(x*h)
	// Note that the edwards25519 package uses Addtion as the group operation
	grgxh := new(curve.Point).Add(gr, gxh)

	return V.Equal(grgxh) == 1
}

func combineErr(e1, e2 error) error {
	if e1 == nil && e2 == nil {
		return nil
	} else if e1 != nil {
		if e2 == nil {
			return e1
		}
		return fmt.Errorf("%v and %v", e1, e2)
	}
	return e2
}

// Verify checks for both, ProofX and ProofY that
// TODO
func (v *Vote) VerifyProofs() (ok bool) {
	okX := verifyProof(&v.ProofX.PointV, &v.ProofX.ScalarR, &v.privX, v.Index)
	okR := verifyProof(&v.ProofR.PointV, &v.ProofR.ScalarR, &v.privR, v.Index)
	return okX && okR
}

func newVoteWithRand(bit bool, index uint64, rand io.Reader) (vote *Vote, e error) {
	vote = &Vote{
		bit: bit,
	}
	vote.Commitment.Index = index

	e = newPriv(&vote.privX, rand)
	if e != nil {
		return nil, e
	}
	e = newPriv(&vote.privR, rand)
	if e != nil {
		return nil, e
	}

	setPub(&vote.privX, &vote.Commitment.PubX)
	setPub(&vote.privR, &vote.Commitment.PubR)

	e = vote.genProofs()

	return vote, nil
}

func NewVote(bit bool, index uint64) (vote *Vote, e error) {
	return newVoteWithRand(bit, index, nil)
}

func pubStr(p *curve.Point) string {
	return b32.EncodeToString(p.Bytes())
}

func (c *Commitment) String() string {
	return fmt.Sprintf(`{"PubX": "%s", "PubR": "%s"}`, pubStr(&c.PubX), pubStr(&c.PubR))
}

func (c *Commitment) MarshalJSON() ([]byte, error) {
	s := c.String()
	return []byte(s), nil
}