package stm

import (
	"fmt"
)

// Parameters holds the three scalar values from a Mithril cert's
// protocol_parameters field. Only the three that matter for verification.
type Parameters struct {
	K     uint64  // minimum distinct lottery wins for a valid aggregate
	M     uint64  // total lottery slots per signing round
	PhiF  float64 // lottery success base rate (the "phi_f" param)
}

// Verify runs the full STM aggregate-signature verification:
//
//  1. k-threshold:   DistinctWins(multi_sig) >= params.K
//  2. Lottery check: for each (signer, index), ev < threshold(stake)
//  3. Merkle proof:  each claimed signer leaf is in avk at its index
//  4. BLS verify:    MuSig-aggregated (sig, vk) over (msg || avk.root)
//
// All four must pass. Returns a descriptive error at the first failure.
//
// msg is the ASCII bytes of the certificate's signed_message hex string.
func Verify(msg []byte, ms *MultiSig, avk *AVK, params Parameters) error {
	// (1) k-threshold
	distinct := ms.DistinctWins()
	if uint64(len(distinct)) < params.K {
		return fmt.Errorf("k-threshold: got %d distinct wins, want >= %d", len(distinct), params.K)
	}
	// Also: every claimed index must appear exactly once across all signers.
	// Rust enforces `nr_indices == unique_indices.len()` via HashSet.
	total := ms.TotalWins()
	if total != len(distinct) {
		return fmt.Errorf("lottery indices not unique across signers: total=%d distinct=%d",
			total, len(distinct))
	}

	// Compute msgp = msg || avk.MerkleRoot — used by lottery check AND BLS.
	msgp := append(append([]byte(nil), msg...), avk.MerkleRoot...)

	// (2) Lottery check — per (signer, claimed index).
	for i, s := range ms.Signatures {
		for _, idx := range s.Sig.Indexes {
			if idx >= params.M {
				return fmt.Errorf("signer[%d] claimed index %d >= m=%d", i, idx, params.M)
			}
			ev := EvaluateSigma(msgp, idx, s.Sig.Sigma)
			if !IsLotteryWon(params.PhiF, ev, s.RegParty.Stake, avk.TotalStake) {
				return fmt.Errorf("signer[%d] lottery loss at index %d (stake=%d/%d)",
					i, idx, s.RegParty.Stake, avk.TotalStake)
			}
		}
	}

	// (3) Merkle batch proof — prove each signer leaf is in the AVK tree.
	leaves := make([][]byte, len(ms.Signatures))
	indices := make([]uint64, len(ms.Signatures))
	// Rust sorts leaves by signer_index ascending; so must we.
	sorted := make([]int, len(ms.Signatures))
	for i := range sorted {
		sorted[i] = i
	}
	// Sort indices ascending while tracking permutation
	for i := 0; i < len(sorted); i++ {
		for j := i + 1; j < len(sorted); j++ {
			if ms.Signatures[sorted[j]].Sig.SignerIndex < ms.Signatures[sorted[i]].Sig.SignerIndex {
				sorted[i], sorted[j] = sorted[j], sorted[i]
			}
		}
	}
	for outIdx, origIdx := range sorted {
		s := ms.Signatures[origIdx]
		leaves[outIdx] = LeafBytes(s.RegParty.VK, s.RegParty.Stake)
		indices[outIdx] = s.Sig.SignerIndex
	}
	proofVals := make([][]byte, len(ms.BatchProof.Values))
	for i, v := range ms.BatchProof.Values {
		proofVals[i] = v
	}
	if err := VerifyMerkleBatch(avk.MerkleRoot, int(avk.NumLeaves), leaves, indices, proofVals); err != nil {
		return fmt.Errorf("merkle batch proof: %w", err)
	}

	// (4) BLS aggregate verify — unique signers, no multiplicity.
	sigs := make([][]byte, len(ms.Signatures))
	vks := make([][]byte, len(ms.Signatures))
	for i, s := range ms.Signatures {
		sigs[i] = s.Sig.Sigma
		vks[i] = s.RegParty.VK
	}
	if err := BlsAggregateVerify(msgp, sigs, vks); err != nil {
		return fmt.Errorf("bls aggregate: %w", err)
	}

	return nil
}