package stm

import (
	"encoding/binary"
	"math"
	"math/big"

	"golang.org/x/crypto/blake2b"
)

// EvaluateSigma computes the 64-byte lottery evaluation for a given
// (msg, index, sigma). Mirrors Rust's evaluate_dense_mapping:
//
//	ev = Blake2b-512( "map" || msg || le_u64(index) || sigma_bytes )
//
// The 64-byte output is the lottery draw, interpreted as a big unsigned
// integer in LSF/little-endian byte order per the Rust impl:
//
//	rug::Integer::from_digits(&ev, Order::LsfLe)
//	num_bigint::BigInt::from_bytes_le(Sign::Plus, &ev)
func EvaluateSigma(msg []byte, index uint64, sigma []byte) [64]byte {
	h, _ := blake2b.New512(nil)
	h.Write([]byte("map"))
	h.Write(msg)
	var idxBuf [8]byte
	binary.LittleEndian.PutUint64(idxBuf[:], index)
	h.Write(idxBuf[:])
	h.Write(sigma)
	var out [64]byte
	copy(out[:], h.Sum(nil))
	return out
}

// evAsBigInt converts the 64-byte ev output to a big.Int using LE byte
// order (matching the Rust `from_bytes_le`).
func evAsBigInt(ev [64]byte) *big.Int {
	// big.Int.SetBytes is BE; so flip.
	rev := make([]byte, len(ev))
	for i := range ev {
		rev[i] = ev[len(ev)-1-i]
	}
	return new(big.Int).SetBytes(rev)
}

// IsLotteryWon reports whether a signer with the given stake wins the
// lottery at the claimed index for the given ev.
//
// Predicate:  p < 1 - (1 - phi_f)^w,  where
//
//	p     = ev / 2^512
//	w     = stake / total_stake
//	phi_f = protocol parameter in (0, 1]
//
// Equivalent reformulation (used here): `q < exp(-w * c)` where
// `q = 1/(1-p)` and `c = ln(1 - phi_f)`. Evaluated via Taylor series
// with early-stop on the error bound (constant M=3 from the Rust impl).
func IsLotteryWon(phiF float64, ev [64]byte, stake, totalStake uint64) bool {
	if math.Abs(phiF-1.0) < 1e-15 {
		return true
	}
	// Defensive: zero-stake or zero-total-stake produces nonsense (and
	// totalStake==0 would panic at SetFrac). Guard at the lottery layer
	// in addition to AVK-decode-time validation.
	if stake == 0 || totalStake == 0 {
		return false
	}

	// ev as big int (LE interpretation)
	evInt := evAsBigInt(ev)

	// evMax = 2^512
	evMax := new(big.Int).Lsh(big.NewInt(1), 512)

	// q = evMax / (evMax - ev)  — a Ratio
	denom := new(big.Int).Sub(evMax, evInt)
	q := new(big.Rat).SetFrac(new(big.Int).Set(evMax), denom)

	// c = ln(1 - phi_f); x = -w * c
	cFloat := math.Log(1.0 - phiF)
	c := new(big.Rat).SetFloat64(cFloat)

	w := new(big.Rat).SetFrac(
		new(big.Int).SetUint64(stake),
		new(big.Int).SetUint64(totalStake),
	)

	x := new(big.Rat).Mul(w, c)
	x.Neg(x)

	return taylorCompare(1000, q, x)
}

// taylorCompare reports whether cmp < exp(x), using a Taylor series
// expansion with an early-stop error heuristic (M = 3).
func taylorCompare(bound int, cmp, x *big.Rat) bool {
	newX := new(big.Rat).Set(x)
	phi := new(big.Rat).SetInt64(1)
	divisor := big.NewInt(1)
	three := big.NewRat(3, 1)
	absNewX := new(big.Rat)
	errorTerm := new(big.Rat)
	sum := new(big.Rat)
	diff := new(big.Rat)

	for i := 0; i < bound; i++ {
		phi.Add(phi, newX)
		divisor = new(big.Int).Add(divisor, big.NewInt(1))
		// newX = newX * x / divisor
		nx := new(big.Rat).Mul(newX, x)
		nx.Quo(nx, new(big.Rat).SetInt(divisor))
		newX = nx

		absNewX.Abs(newX)
		errorTerm.Mul(absNewX, three)

		sum.Add(phi, errorTerm)
		if cmp.Cmp(sum) > 0 {
			return false
		}
		diff.Sub(phi, errorTerm)
		if cmp.Cmp(diff) < 0 {
			return true
		}
	}
	return false
}