package gnarlring

import (
	"crypto/rand"
	"io"
)

// Vote represents a binary choice encoded as a short vectors norm direction.
// A YES vote produces a short vector z where the first coefficient is +1
// (or more generally, has positive aggregate norm). A NO vote produces z
// with negative-norm first coefficient.
type Vote struct {
	Z *Poly27 // short vector encoding the vote
}

// CastVote produces a short-vector vote. sign = +1 for YES, -1 for NO.
func CastVote(sign int) *Vote {
	return CastVoteFrom(sign, nil)
}

// CastVoteFrom produces a vote with a given randomness source.
// The vote is a fresh Gaussian vector with the first coefficient encoding
// the sign: 1 for YES, Q-1 for NO (centered = -1).
func CastVoteFrom(sign int, rng io.Reader) *Vote {
	if rng == nil {
		rng = rand.Reader
	}
	gs := NewGaussSamplerFrom(DefaultSigma(), rng)
	z := gs.SamplePoly()

	if sign >= 0 {
		z.Coeffs[0] = 1
	} else {
		z.Coeffs[0] = Q - 1 // centered = -1
	}
	return &Vote{Z: z}
}

// VoteTally aggregates multiple votes into a single short vector.
// z_total = Σ vote_i.Z. The norm of z_total encodes the net YES-NO count.
type VoteTally struct {
	Z        *Poly27
	YesCount int
	NoCount  int
}

// NewVoteTally creates an empty tally.
func NewVoteTally() *VoteTally {
	return &VoteTally{Z: NewPoly27()}
}

// Add includes a vote in the tally.
func (vt *VoteTally) Add(v *Vote, isYes bool) {
	vt.Z = Add(vt.Z, v.Z)
	if isYes {
		vt.YesCount++
	} else {
		vt.NoCount++
	}
}

// ConsensusResult returns whether the tally meets a k-of-n threshold.
// True if net YES votes ≥ threshold.
func (vt *VoteTally) ConsensusResult(threshold int) bool {
	return vt.YesCount >= threshold
}

// NormEstimate returns the expected norm range for k YES votes among
// n total participants. Each vote contributes approximately
// sqrt(N) * DefaultSigma() ≈ 54 to the squared norm.
func VoteNormEstimate(k int) float64 {
	sigma := DefaultSigma()
	perVote := float64(N) * sigma * sigma // expected ||z||² per vote
	return perVote * float64(k)
}

// VerifyVoteTally checks that the tally's Z vector has norm consistent
// with the expected range for the given yes/no counts.
func VerifyVoteTally(vt *VoteTally) bool {
	expectedYes := VoteNormEstimate(vt.YesCount)
	expectedNo := VoteNormEstimate(vt.NoCount)
	// Net norm: yes contributions mostly cancel with no contributions.
	// The residual norm is sqrt(yes + no) * per_vote * (yes-no)/total.
	// Simplified: check that the observed norm is plausible.
	observed := float64(NormSq(vt.Z))

	// Maximum possible: all votes aligning (yes - no all same sign).
	maxPossible := expectedYes + expectedNo

	// Minimum: complete cancellation.
	minPossible := 0.0

	if observed > maxPossible*1.5 || observed < minPossible {
		return false
	}
	return true
}

// EncryptedVote wraps a vote in Ring-LWE encryption for private broadcast.
type EncryptedVote struct {
	Ct *LWECiphertext // encrypted vote vector (bit-by-bit or as ring element)
}

// EncryptVote encrypts a child's vote under the coordinator's LWE public key.
// The vote's sign is encoded as a single LWE-encrypted bit.
func EncryptVote(lwePK *LWEPublicKey, v *Vote, rng io.Reader) *EncryptedVote {
	if rng == nil {
		rng = rand.Reader
	}
	// Decode the sign from the vote vector: compare centered signs.
	// First coefficient encodes the vote: 1 = YES, Q-1 = NO.
	first := v.Z.Coeffs[0]
	bit := 0
	if first == 1 {
		bit = 1 // centered positive = YES
	}
	ct := LWEEncryptFrom(lwePK, bit, rng)
	return &EncryptedVote{Ct: ct}
}

// DecryptVote decrypts an encrypted vote and returns whether it was YES.
func DecryptVote(lweSK *LWESecretKey, ev *EncryptedVote) bool {
	return LWEDecrypt(lweSK, ev.Ct) == 1
}