package stm

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

	bls12381 "github.com/consensys/gnark-crypto/ecc/bls12-381"
	"golang.org/x/crypto/blake2b"
)

// AggregateBLS performs Mithril's MuSig-style weighted aggregation.
//
// For each signer i with signature σ_i and verification key vk_i:
//
//	h₀ = Blake2b-128(σ_0 ‖ σ_1 ‖ … ‖ σ_{n-1})
//	t_i = Blake2b-128(h₀ ‖ i_be_u64)           // 128-bit scalar
//	aggr_σ = Σ t_i · σ_i                        // in G1
//	aggr_vk = Σ t_i · vk_i                      // in G2
//
// Matches mithril-stm::bls_multi_signature::signature::BlsSignature::aggregate.
// The per-signer scalar prevents rogue-key attacks where an adversary chooses
// their vk to cancel the contribution of honest signers.
//
// Single-signer short-circuit: when n == 1 the function returns the raw
// sig/vk unchanged, matching blst's `if vks.len() < 2 { return (vks[0], sigs[0]) }`.
func AggregateBLS(sigs, vks [][]byte) (aggSig, aggVK []byte, err error) {
	if len(sigs) == 0 || len(sigs) != len(vks) {
		return nil, nil, fmt.Errorf("sig/vk mismatch: %d/%d", len(sigs), len(vks))
	}
	if len(sigs) == 1 {
		// No aggregation needed; blst returns the single contribution as-is.
		return sigs[0], vks[0], nil
	}

	// Decode all points + subgroup-check up front.
	sigPts := make([]bls12381.G1Affine, len(sigs))
	vkPts := make([]bls12381.G2Affine, len(vks))
	for i := range sigs {
		if len(sigs[i]) != bls12381.SizeOfG1AffineCompressed {
			return nil, nil, fmt.Errorf("sig[%d] size %d", i, len(sigs[i]))
		}
		if _, err := sigPts[i].SetBytes(sigs[i]); err != nil {
			return nil, nil, fmt.Errorf("sig[%d] decode: %w", i, err)
		}
		if !sigPts[i].IsInSubGroup() {
			return nil, nil, fmt.Errorf("sig[%d] not in subgroup", i)
		}
		if len(vks[i]) != bls12381.SizeOfG2AffineCompressed {
			return nil, nil, fmt.Errorf("vk[%d] size %d", i, len(vks[i]))
		}
		if _, err := vkPts[i].SetBytes(vks[i]); err != nil {
			return nil, nil, fmt.Errorf("vk[%d] decode: %w", i, err)
		}
		if !vkPts[i].IsInSubGroup() {
			return nil, nil, fmt.Errorf("vk[%d] not in subgroup", i)
		}
	}

	// Step 1: h₀ = Blake2b-128(σ_0 ‖ σ_1 ‖ … ‖ σ_{n-1})
	//
	// Note: the Rust impl updates the hasher with `sig.to_bytes()`, which is
	// the compressed 48-byte encoding. We use the input bytes directly —
	// they're the same thing, the signer ships compressed form on the wire.
	base, err := blake2b.New(16, nil) // 128-bit output
	if err != nil {
		return nil, nil, err
	}
	for _, s := range sigs {
		base.Write(s)
	}
	baseDigest := base.Sum(nil) // 16 bytes

	// Step 2: scalars t_i = Blake2b-128(h₀ ‖ be_u64(i))
	//
	// Critical detail: the Rust uses `index.to_be_bytes()` where `index` is
	// `usize`. On 64-bit platforms that's 8 BE bytes. blst_lib then treats
	// the resulting 128-bit scalar as little-endian when doing the mult.
	scalars := make([]*big.Int, len(sigs))
	idxBuf := make([]byte, 8)
	for i := range sigs {
		h, err := blake2b.New(16, nil)
		if err != nil {
			return nil, nil, err
		}
		h.Write(baseDigest)
		binary.BigEndian.PutUint64(idxBuf, uint64(i))
		h.Write(idxBuf)
		digest := h.Sum(nil) // 16 bytes LE
		// blst interprets the scalar blob as little-endian; gnark's
		// ScalarMultiplication takes a big.Int (treated as absolute
		// value mod r). Interpret as LE to match blst.
		scalars[i] = new(big.Int).SetBytes(reverseBytes(digest))
	}

	// Step 3: aggr_sig = Σ t_i · σ_i
	var aggSigJac bls12381.G1Jac
	for i := range sigPts {
		var scaled bls12381.G1Affine
		scaled.ScalarMultiplication(&sigPts[i], scalars[i])
		var jac bls12381.G1Jac
		jac.FromAffine(&scaled)
		aggSigJac.AddAssign(&jac)
	}
	var aggSigAff bls12381.G1Affine
	aggSigAff.FromJacobian(&aggSigJac)
	sigBytes := aggSigAff.Bytes()

	// aggr_vk = Σ t_i · vk_i
	var aggVKJac bls12381.G2Jac
	for i := range vkPts {
		var scaled bls12381.G2Affine
		scaled.ScalarMultiplication(&vkPts[i], scalars[i])
		var jac bls12381.G2Jac
		jac.FromAffine(&scaled)
		aggVKJac.AddAssign(&jac)
	}
	var aggVKAff bls12381.G2Affine
	aggVKAff.FromJacobian(&aggVKJac)
	vkBytes := aggVKAff.Bytes()

	return sigBytes[:], vkBytes[:], nil
}

func reverseBytes(b []byte) []byte {
	out := make([]byte, len(b))
	for i, v := range b {
		out[len(b)-1-i] = v
	}
	return out
}

// BlsAggregateVerify aggregates per-signer contributions (with the
// Mithril MuSig-style scalar weighting) and does one pairing check.
// Does NOT check lottery wins, Merkle membership, or the k-threshold —
// those are separate STM-layer checks.
func BlsAggregateVerify(msg []byte, sigs, vks [][]byte) error {
	aggSig, aggVK, err := AggregateBLS(sigs, vks)
	if err != nil {
		return fmt.Errorf("aggregate: %w", err)
	}
	return BlsVerify(msg, aggSig, aggVK)
}