package veto

import (
	"bytes"
	"crypto/sha512"
	"encoding/json"
	"fmt"
	"io"
	"slices"

	"kesim.org/seal/curve"
)

type Scalar = curve.Curve25519Scalar
type Point = curve.Curve25519Point

var Curve = curve.Curve25519

// Representation of a vote with veto (if set to true)
type Vote struct {
	veto    bool
	private struct {
		id *Scalar
		x  *Scalar
		r  *Scalar
	}
	com *Commitment
}

// Commitment represents the public data sent by a participant
// in round 1 of the protocol.  It is generated out of a Vote.
type Commitment struct {
	Id     *Point `json:"identity"`
	Points struct {
		X *Point
		R *Point
	} `json:"points"`
	Proofs struct {
		X *Proof
		R *Proof
	} `json:"proofs"`
}

// A Schnorr signature to prove knowledge of v for given g^v 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)
type Proof struct {
	V *Point  `json:"V"`
	R *Scalar `json:"r"`
}

// 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 proof(x *Scalar, id *Point) (pr *Proof, e error) {
	pr = &Proof{}

	// choose random v
	v, e := Curve.ScalarFromReader(nil)
	if e != nil {
		return nil, e
	}

	// calculate g^v
	pr.V = Curve.Exp(v)

	// calculate g^x
	gx := Curve.Exp(x)

	// calculate h := H(g, g^v, g^x, i)
	h, e := hash(pr.V, gx, id)
	if e != nil {
		return nil, e
	}

	// Calculate r := v - x*h
	xh := x.Mul(h)
	r := v.Sub(xh)
	pr.R = r

	return pr, nil
}

// Calculate  h := H(g, g^v, g^x, i)
func hash(gv, gx *Point, id *Point) (*Scalar, error) {
	h512 := sha512.New()
	h512.Write(Curve.Identity().Bytes())
	h512.Write(gv.Bytes())
	h512.Write(gx.Bytes())
	h512.Write(id.Bytes())
	hb := h512.Sum(nil)
	return Curve.ScalarFromBytes(hb)
}

func combineErr(es ...error) error {
	var re error
	for _, e := range es {
		if e != nil {
			if re == nil {
				re = e
			} else {
				re = fmt.Errorf("%v and %v", re, e)
			}
		}
	}
	return re
}

// Verifies that g^v == g^r*g^(x*h)
func verifyProof(V *Point, Gx *Point, r *Scalar, id *Point) (ok bool) {
	// Calculate h = H(g, g^v, g^x, id)
	h, e := hash(V, Gx, id)
	if e != nil {
		return false
	}

	// Calculate g^(x*h) = (g^x)^h
	gxh := Gx.Exp(h)

	// Calculate g^r
	gr := Curve.Exp(r)

	// Calculate g^r*g^(x*h)
	grgxh := gr.Mul(gxh)

	// Return true if g^v == g^r*g^(x*h)
	return V.Equal(grgxh)
}

// Verify verifies the proofs for both, g^x and g^r
func (c *Commitment) VerifyProofs() (ok bool) {
	okX := verifyProof(c.Proofs.X.V, c.Points.X, c.Proofs.X.R, c.Id)
	okR := verifyProof(c.Proofs.R.V, c.Points.R, c.Proofs.R.R, c.Id)
	return okX && okR
}

// Generates a vote with commitments and proofs and takes the input for
// the randomness from the given io.Reader
func newVoteWithRand(veto bool, rand io.Reader) (v *Vote, e error) {
	v = &Vote{
		veto: veto,
	}

	var e1, e2, e3 error

	v.private.id, e1 = Curve.ScalarFromReader(rand)
	v.private.x, e2 = Curve.ScalarFromReader(rand)
	v.private.r, e3 = Curve.ScalarFromReader(rand)

	e = combineErr(e1, e2, e3)
	if e != nil {
		return nil, e
	}

	c := new(Commitment)
	v.com = c
	c.Id = Curve.Exp(v.private.id)
	c.Points.X = Curve.Exp(v.private.x)
	c.Points.R = Curve.Exp(v.private.r)
	c.Proofs.X, e1 = proof(v.private.x, c.Id)
	c.Proofs.R, e2 = proof(v.private.r, c.Id)

	return v, combineErr(e1, e2)
}

// NewVote generates a vote for given bit and index, taking crypt/Reader as
// source for randomness
func NewVote(bit bool) (vote *Vote, e error) {
	return newVoteWithRand(bit, nil)
}

// Generate the commitment to the Vote
func (v *Vote) Commitment() *Commitment {
	return v.com
}

func (c *Commitment) String() string {
	buf := &bytes.Buffer{}
	dec := json.NewEncoder(buf)
	dec.SetIndent("", "  ")
	e := dec.Encode(c)
	if e != nil {
		return fmt.Sprintf("<error encoding: %v>", e)
	}
	return buf.String()
}

type coms []*Commitment

func (coms coms) prod() (product *Point) {
	product = Curve.Identity()
	for _, com := range coms {
		product = product.Mul(com.Points.X)
	}
	return product
}

// For a given slice of commitments of length n, compute
// g^y_i = Prod(g^x_j, 1, i-1)/Prod(g^x_j, i+1, n)
func (coms coms) computeGy(index int) *Point {
	gy1 := coms[:index].prod()
	gy2 := coms[index+1:].prod().Inv()
	return gy1.Mul(gy2)
}

// Round2 implements the round 2 of the AV-Net protocol, where a participant
// has received all commitments from the other participants and calculates
// g^(c*y), where c is either private.r (when veto is true) or private.x.
func (v *Vote) Round2(participants []*Commitment) (gcy *Point, e error) {
	var c int
	for _, p := range participants {
		if p.Id.Equal(v.com.Id) {
			c += 1
		}
	}
	if c == 0 {
		return nil, fmt.Errorf("/me(%s) not in participants", v.com.Id)
	} else if c > 1 {
		return nil, fmt.Errorf("/me(%s) %d times in participants", v.com.Id, c)
	}

	slices.SortStableFunc(participants, func(a, b *Commitment) int {
		return bytes.Compare(a.Id.Bytes(), b.Id.Bytes())
	})

	index := slices.IndexFunc(participants, func(c *Commitment) bool { return c.Id.Equal(v.com.Id) })
	if index < 0 {
		// Should not happen!
		return nil, fmt.Errorf("Wait, what!? Couldn't find /me(%s) after sorting!", v.com.Id)
	}
	gy := (coms(participants)).computeGy(index)

	if v.veto {
		return gy.Exp(v.private.r), nil
	}
	return gy.Exp(v.private.x), nil
}

type points []*Point

// IsVetoed is the final step in the AV-Net protocol, where each participant, has
// received the g^(c_i*y_i) from all other participants and calculates the product
// of them.  If the result is the unit element of the group, no veto was present.
func (pts points) IsVetoed() bool {
	product := Curve.Product(pts...)
	one := Curve.Identity()
	return !one.Equal(product)
}