// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package thresholdvm

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/luxfi/log"
	"github.com/luxfi/threshold/pkg/math/curve"
	"github.com/luxfi/threshold/pkg/party"
	"github.com/luxfi/threshold/pkg/pool"
	"github.com/luxfi/threshold/pkg/protocol"
	"github.com/luxfi/threshold/protocols/bls"
	"github.com/luxfi/threshold/protocols/cmp"
	"github.com/luxfi/threshold/protocols/frost"
	"github.com/luxfi/threshold/protocols/lss"
	"github.com/luxfi/threshold/protocols/corona"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// =============================================================================
// Test Harness - Local implementation for MPC protocol testing
// =============================================================================

// testNetwork simulates a network for MPC protocol testing
type testNetwork struct {
	parties  map[party.ID]chan *protocol.Message
	handlers map[party.ID]*protocol.Handler
	mu       sync.RWMutex
	wg       sync.WaitGroup
	done     chan struct{}
}

func newTestNetwork(ids []party.ID) *testNetwork {
	n := &testNetwork{
		parties:  make(map[party.ID]chan *protocol.Message),
		handlers: make(map[party.ID]*protocol.Handler),
		done:     make(chan struct{}),
	}
	for _, id := range ids {
		n.parties[id] = make(chan *protocol.Message, 10000)
	}
	return n
}

func (n *testNetwork) close() {
	close(n.done)
	n.wg.Wait()
	for _, ch := range n.parties {
		close(ch)
	}
}

// sessionCounter provides unique session IDs for each protocol run
var sessionCounter uint64

// runTestProtocol runs a protocol to completion across all parties
func runTestProtocol(
	t testing.TB,
	ids []party.ID,
	createStart func(id party.ID) protocol.StartFunc,
) (map[party.ID]interface{}, error) {
	network := newTestNetwork(ids)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) // Longer timeout for slow protocols
	defer cancel()
	defer network.close()

	// Use real logger to see protocol errors
	logger := log.NewTestLogger(log.DebugLevel)
	results := sync.Map{}
	errors := sync.Map{}

	// Use proper protocol config with explicit timeouts - matching threshold library harness
	config := &protocol.Config{
		Workers:         4,
		PriorityWorkers: 4,
		BufferSize:      10000,
		PriorityBuffer:  1000,
		MessageTimeout:  30 * time.Second,
		RoundTimeout:    60 * time.Second,
		ProtocolTimeout: 5 * time.Minute, // Match harness - don't let handler create its own timeout
	}

	// Generate unique session ID for this protocol run
	sessionCounter++
	sessionID := []byte(fmt.Sprintf("test-session-%d-%d", time.Now().UnixNano(), sessionCounter))

	// Create handlers for all parties
	// CRITICAL: Use context.Background() for handlers, not the timeout context!
	// The harness manages timeouts externally. Passing a timeout context to NewHandler
	// can cause premature cancellation of handler internal operations.
	for _, id := range ids {
		startFunc := createStart(id)
		handler, err := protocol.NewHandler(
			context.Background(), // Don't use ctx to avoid premature cancellation
			logger,
			nil, // No metrics registry for tests
			startFunc,
			sessionID,
			config,
		)
		if err != nil {
			return nil, err
		}
		network.handlers[id] = handler
	}

	// Start message routing for each party
	for _, partyID := range ids {
		id := partyID
		handler := network.handlers[id]

		// Outgoing messages
		network.wg.Add(1)
		go func() {
			defer network.wg.Done()
			for {
				select {
				case <-network.done:
					return
				case msg := <-handler.Listen():
					if msg == nil {
						return
					}
					// Route message - collect targets first, then send outside lock
					network.mu.RLock()
					targets := make([]chan *protocol.Message, 0)
					if msg.To == "" {
						// Broadcast to all parties except sender
						for toID, ch := range network.parties {
							if toID != id {
								targets = append(targets, ch)
							}
						}
					} else {
						// Point-to-point
						if ch, ok := network.parties[party.ID(msg.To)]; ok {
							targets = append(targets, ch)
						}
					}
					network.mu.RUnlock()

					// Send outside lock to avoid deadlock - MUST deliver, don't drop!
					for _, ch := range targets {
						ch <- msg
					}
				}
			}
		}()

		// Incoming messages
		network.wg.Add(1)
		go func() {
			defer network.wg.Done()
			ch := network.parties[id]
			for {
				select {
				case <-network.done:
					return
				case msg := <-ch:
					if msg != nil {
						handler.Accept(msg)
					}
				}
			}
		}()
	}

	// Wait for all parties to complete with timeout
	var wg sync.WaitGroup
	for id, handler := range network.handlers {
		id := id
		handler := handler
		wg.Add(1)
		go func() {
			defer wg.Done()
			result, err := handler.WaitForResult()
			if err != nil {
				errors.Store(id, err)
			} else {
				results.Store(id, result)
			}
		}()
	}

	// Wait with external timeout context
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		// All handlers completed
	case <-ctx.Done():
		return nil, fmt.Errorf("protocol timed out after 10 minutes")
	}

	// Check for errors - log all errors for debugging
	var errs []error
	errors.Range(func(key, value interface{}) bool {
		t.Logf("Handler %v error: %v", key, value)
		errs = append(errs, value.(error))
		return true
	})
	if len(errs) > 0 {
		return nil, errs[0]
	}

	// Collect results
	resultMap := make(map[party.ID]interface{})
	results.Range(func(key, value interface{}) bool {
		resultMap[key.(party.ID)] = value
		return true
	})

	return resultMap, nil
}

// testPartyIDs generates party IDs for testing
func testPartyIDs(n int) []party.ID {
	ids := make([]party.ID, n)
	for i := 0; i < n; i++ {
		ids[i] = party.ID(string(rune('a' + i)))
	}
	return ids
}

// =============================================================================
// Unit Tests
// =============================================================================

// TestProtocolExecutorCreation tests creating a ProtocolExecutor
func TestProtocolExecutorCreation(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)

	pe := NewProtocolExecutor(workerPool, logger)
	require.NotNil(pe)
	require.NotNil(pe.pool)
	require.NotNil(pe.handlers)
}

// TestLSSKeygenStartFunc tests LSS key generation start function
func TestLSSKeygenStartFunc(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	pIDs := []party.ID{"alice", "bob", "charlie"}
	selfID := party.ID("alice")
	threshold := 2

	startFunc := pe.LSSKeygenStartFunc(selfID, pIDs, threshold)
	require.NotNil(startFunc)
}

// TestCMPKeygenStartFunc tests CMP key generation start function
func TestCMPKeygenStartFunc(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	pIDs := []party.ID{"alice", "bob", "charlie"}
	selfID := party.ID("alice")
	threshold := 2

	startFunc := pe.CMPKeygenStartFunc(selfID, pIDs, threshold)
	require.NotNil(startFunc)
}

// TestFROSTKeygenStartFunc tests FROST key generation start function
func TestFROSTKeygenStartFunc(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	pIDs := []party.ID{"alice", "bob", "charlie"}
	selfID := party.ID("alice")
	threshold := 2

	startFunc := pe.FROSTKeygenStartFunc(selfID, pIDs, threshold)
	require.NotNil(startFunc)
}

// TestFROSTKeygenTaprootStartFunc tests FROST Taproot key generation
func TestFROSTKeygenTaprootStartFunc(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	pIDs := []party.ID{"alice", "bob", "charlie"}
	selfID := party.ID("alice")
	threshold := 2

	startFunc := pe.FROSTKeygenTaprootStartFunc(selfID, pIDs, threshold)
	require.NotNil(startFunc)
}

// TestHandlerLifecycle tests creating, getting, and removing handlers
func TestHandlerLifecycle(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	pIDs := []party.ID{"alice", "bob", "charlie"}
	selfID := party.ID("alice")
	threshold := 2

	ctx := context.Background()
	sessionID := "test-session-1"
	startFunc := pe.FROSTKeygenStartFunc(selfID, pIDs, threshold)

	handler, err := pe.CreateHandler(ctx, sessionID, startFunc)
	require.NoError(err)
	require.NotNil(handler)

	retrieved, ok := pe.GetHandler(sessionID)
	require.True(ok)
	require.Equal(handler, retrieved)

	_, ok = pe.GetHandler("non-existent")
	require.False(ok)

	pe.RemoveHandler(sessionID)

	_, ok = pe.GetHandler(sessionID)
	require.False(ok)
}

// TestKeyShareWrappers tests the KeyShare wrapper implementations
func TestKeyShareWrappers(t *testing.T) {
	require := require.New(t)

	lssShare := &LSSKeyShare{Config: nil}
	require.Equal(ProtocolLSS, lssShare.Protocol())

	cmpShare := &CMPKeyShare{Config: nil}
	require.Equal(ProtocolCGGMP21, cmpShare.Protocol())

	frostShare := &FROSTKeyShare{Config: nil}
	require.Equal(ProtocolFrost, frostShare.Protocol())
}

// TestECDSASignatureType tests ECDSA signature type
func TestECDSASignatureType(t *testing.T) {
	require := require.New(t)

	sig := &ECDSASignature{
		R: make([]byte, 32),
		S: make([]byte, 32),
		V: 27,
	}

	require.Len(sig.R, 32)
	require.Len(sig.S, 32)
	require.Equal(byte(27), sig.V)
}

// TestSchnorrSignatureType tests Schnorr signature type
func TestSchnorrSignatureType(t *testing.T) {
	require := require.New(t)

	sig := &SchnorrSignature{
		R: make([]byte, 32),
		Z: make([]byte, 32),
	}

	require.Len(sig.R, 32)
	require.Len(sig.Z, 32)
}

// =============================================================================
// Full Protocol Execution Tests
// =============================================================================

// TestFROSTKeygenFullExecution runs a complete FROST keygen
func TestFROSTKeygenFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping FROST keygen execution in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(5)
	threshold := 3

	results, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return frost.Keygen(curve.Secp256k1{}, id, pIDs, threshold)
	})
	require.NoError(err)
	require.Len(results, 5)

	var firstPubKey curve.Point
	for id, result := range results {
		config, ok := result.(*frost.Config)
		require.True(ok, "result should be *frost.Config for party %s", id)
		require.NotNil(config)
		require.NotNil(config.PublicKey)

		if firstPubKey == nil {
			firstPubKey = config.PublicKey
		} else {
			assert.True(t, firstPubKey.Equal(config.PublicKey),
				"all parties should have same public key")
		}
	}

	t.Logf("FROST keygen completed: 5 parties, threshold=%d", threshold)
}

// TestLSSKeygenFullExecution runs a complete LSS keygen
func TestLSSKeygenFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping LSS keygen execution in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(3)
	threshold := 2

	// LSS uses a shared pool - use pool.NewPool(0) for default workers
	// This matches the LSS library's own integration tests
	sharedPool := pool.NewPool(0)
	defer sharedPool.TearDown()

	results, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return lss.Keygen(curve.Secp256k1{}, id, pIDs, threshold, sharedPool)
	})
	require.NoError(err)
	require.Len(results, 3)

	for id, result := range results {
		config, ok := result.(*lss.Config)
		require.True(ok, "result should be *lss.Config for party %s", id)
		require.NotNil(config)
	}

	t.Logf("LSS keygen completed: 3 parties, threshold=%d", threshold)
}

// TestCMPKeygenFullExecution runs a complete CMP keygen
func TestCMPKeygenFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping CMP keygen execution in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(3)
	threshold := 2

	// CRITICAL: Each party needs its OWN pool!
	// The Pool type is NOT thread-safe for concurrent Search calls.
	// Sharing a pool across parties causes deadlock when multiple parties
	// concurrently call Finalize which uses pool.Search for Paillier prime generation.
	pools := make(map[party.ID]*pool.Pool)
	for _, id := range pIDs {
		pools[id] = pool.NewPool(4)
	}
	defer func() {
		for _, p := range pools {
			p.TearDown()
		}
	}()

	results, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return cmp.Keygen(curve.Secp256k1{}, id, pIDs, threshold, pools[id])
	})
	require.NoError(err)
	require.Len(results, 3)

	var firstPubKey curve.Point
	for id, result := range results {
		config, ok := result.(*cmp.Config)
		require.True(ok, "result should be *cmp.Config for party %s", id)
		require.NotNil(config)

		pubKey := config.PublicPoint()
		if firstPubKey == nil {
			firstPubKey = pubKey
		} else {
			assert.True(t, firstPubKey.Equal(pubKey),
				"all parties should have same public key")
		}
	}

	t.Logf("CMP keygen completed: 3 parties, threshold=%d", threshold)
}

// TestFROSTSignFullExecution runs a complete FROST keygen + sign
func TestFROSTSignFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping FROST sign execution in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(5)
	threshold := 3
	// FROST requires threshold+1 signers for a t-of-n threshold signature
	signers := pIDs[:threshold+1]
	message := []byte("test message for threshold signing")

	// Keygen
	keygenResults, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return frost.Keygen(curve.Secp256k1{}, id, pIDs, threshold)
	})
	require.NoError(err)
	require.Len(keygenResults, 5)

	firstConfig := keygenResults[pIDs[0]].(*frost.Config)
	publicKey := firstConfig.PublicKey

	// Sign
	signResults, err := runTestProtocol(t, signers, func(id party.ID) protocol.StartFunc {
		config := keygenResults[id].(*frost.Config)
		return frost.Sign(config, signers, message)
	})
	require.NoError(err)
	require.Len(signResults, len(signers))

	// Verify
	for _, result := range signResults {
		var sig *frost.Signature
		switch s := result.(type) {
		case *frost.Signature:
			sig = s
		case frost.Signature:
			sig = &s
		default:
			t.Fatalf("unexpected signature type: %T", result)
		}
		require.NotNil(sig)
		assert.True(t, sig.Verify(publicKey, message), "signature should verify")
		break
	}

	t.Logf("FROST sign completed: %d signers, %d-of-%d threshold signature verified", len(signers), threshold, len(pIDs))
}

// TestLSSSignFullExecution runs a complete LSS keygen + sign
// The livelock bug in the threshold library has been fixed.
func TestLSSSignFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping LSS sign execution in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(5)
	threshold := 3
	// LSS uses exactly `threshold` signers (not threshold+1 like CMP)
	signers := pIDs[:threshold]
	// LSS sign requires exactly 32 bytes (SHA-256 hash)
	messageHash := make([]byte, 32)
	copy(messageHash, []byte("test message hash for LSS sign"))

	// Create per-party pools for keygen
	keygenPools := make(map[party.ID]*pool.Pool)
	for _, id := range pIDs {
		keygenPools[id] = pool.NewPool(4)
	}
	defer func() {
		for _, p := range keygenPools {
			p.TearDown()
		}
	}()

	// Run LSS Keygen
	t.Log("Running LSS keygen...")
	keygenResults, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return lss.Keygen(curve.Secp256k1{}, id, pIDs, threshold, keygenPools[id])
	})
	require.NoError(err)
	require.Len(keygenResults, len(pIDs))

	// Extract configs
	configs := make(map[party.ID]*lss.Config)
	for id, result := range keygenResults {
		config, ok := result.(*lss.Config)
		require.True(ok, "result should be *lss.Config for party %s", id)
		configs[id] = config
	}
	t.Logf("LSS keygen completed: %d parties, threshold %d", len(pIDs), threshold)

	// Create per-party pools for signing
	signPools := make(map[party.ID]*pool.Pool)
	for _, id := range signers {
		signPools[id] = pool.NewPool(4)
	}
	defer func() {
		for _, p := range signPools {
			p.TearDown()
		}
	}()

	// Run LSS Sign with threshold signers
	t.Logf("Running LSS sign with signers: %v", signers)
	signResults, err := runTestProtocol(t, signers, func(id party.ID) protocol.StartFunc {
		return lss.Sign(configs[id], signers, messageHash, signPools[id])
	})
	require.NoError(err)
	require.Len(signResults, len(signers))

	t.Logf("LSS sign completed: %d signers, %d-of-%d threshold signature", len(signers), threshold, len(pIDs))
}

// TestCMPSignFullExecution runs a complete CMP keygen + sign
func TestCMPSignFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping CMP sign execution in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(3)
	threshold := 2
	// CMP requires threshold+1 parties for signing, use all 3
	signers := pIDs
	message := []byte("test message for CMP threshold signing")

	// CRITICAL: Each party needs its OWN pool!
	// The Pool type is NOT thread-safe for concurrent Search calls.
	keygenPools := make(map[party.ID]*pool.Pool)
	for _, id := range pIDs {
		keygenPools[id] = pool.NewPool(4)
	}
	defer func() {
		for _, p := range keygenPools {
			p.TearDown()
		}
	}()

	// Keygen
	keygenResults, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return cmp.Keygen(curve.Secp256k1{}, id, pIDs, threshold, keygenPools[id])
	})
	require.NoError(err)
	require.Len(keygenResults, 3)

	// Sign with all parties (CMP requires > threshold parties)
	// Create fresh per-party pools for signing
	signPools := make(map[party.ID]*pool.Pool)
	for _, id := range signers {
		signPools[id] = pool.NewPool(4)
	}
	defer func() {
		for _, p := range signPools {
			p.TearDown()
		}
	}()

	signResults, err := runTestProtocol(t, signers, func(id party.ID) protocol.StartFunc {
		config := keygenResults[id].(*cmp.Config)
		return cmp.Sign(config, signers, message, signPools[id])
	})
	require.NoError(err)
	require.Len(signResults, len(signers))

	t.Logf("CMP sign completed: %d-of-%d threshold signature", len(signers), len(pIDs))
}

// TestProtocolExecutorWithRealKeygen tests ProtocolExecutor wrapper with real keygen
func TestProtocolExecutorWithRealKeygen(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in short mode")
	}

	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	pIDs := testPartyIDs(3)
	threshold := 2

	startFunc := pe.LSSKeygenStartFunc(pIDs[0], pIDs, threshold)
	require.NotNil(startFunc)

	results, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return pe.LSSKeygenStartFunc(id, pIDs, threshold)
	})
	require.NoError(err)
	require.Len(results, 3)

	for id, result := range results {
		config, ok := result.(*lss.Config)
		require.True(ok, "result should be *lss.Config for party %s", id)

		share := &LSSKeyShare{Config: config}
		require.Equal(ProtocolLSS, share.Protocol())
		require.NotNil(share.PublicKey())
		require.Equal(id, share.PartyID())
		require.Equal(threshold, share.Threshold())
		require.Equal(3, share.TotalParties())
	}

	t.Logf("ProtocolExecutor+LSSKeyShare integration test passed")
}

// =============================================================================
// BLS Threshold Network Tests
// =============================================================================

// TestBLSThresholdSigningFullExecution tests BLS threshold keygen + signing
// Uses TrustedDealer for key generation and tests threshold signature aggregation
func TestBLSThresholdSigningFullExecution(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping BLS threshold signing in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(5)
	threshold := 3
	message := []byte("test message for BLS threshold signing")

	// Generate keys using TrustedDealer
	dealer := &bls.TrustedDealer{
		Threshold:    threshold,
		TotalParties: len(pIDs),
	}

	shares, groupPK, err := dealer.GenerateShares(context.Background(), pIDs)
	require.NoError(err)
	require.Len(shares, 5)
	require.NotNil(groupPK)

	// Get verification keys
	verificationKeys := bls.GetVerificationKeys(shares)
	require.Len(verificationKeys, 5)

	// Create configs for each party
	configs := make(map[party.ID]*bls.Config)
	for _, id := range pIDs {
		configs[id] = bls.NewConfig(id, threshold, len(pIDs), shares[id], groupPK, verificationKeys)
	}

	// Each party creates a partial signature
	sigShares := make([]*bls.SignatureShare, 0, threshold)
	for i := 0; i < threshold; i++ {
		id := pIDs[i]
		config := configs[id]

		share, err := config.Sign(message)
		require.NoError(err)
		require.NotNil(share)

		// Verify partial signature
		valid := config.VerifyPartialSignature(share, message)
		assert.True(t, valid, "partial signature from %s should be valid", id)

		sigShares = append(sigShares, share)
	}

	// Aggregate signatures
	aggregatedSig, err := bls.AggregateSignatures(sigShares, threshold)
	require.NoError(err)
	require.NotNil(aggregatedSig)

	// Verify aggregated signature against group public key
	valid := configs[pIDs[0]].VerifyAggregateSignature(message, aggregatedSig)
	assert.True(t, valid, "aggregated signature should verify against group public key")

	t.Logf("BLS threshold signing completed: %d-of-%d threshold signature verified", threshold, len(pIDs))
}

// TestBLSThresholdWithDifferentSignerSets tests that any t-of-n signers can produce valid signature
func TestBLSThresholdWithDifferentSignerSets(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping BLS signer set test in short mode")
	}

	require := require.New(t)

	pIDs := testPartyIDs(5)
	threshold := 3
	message := []byte("test message for different signer sets")

	// Generate keys
	dealer := &bls.TrustedDealer{
		Threshold:    threshold,
		TotalParties: len(pIDs),
	}

	shares, groupPK, err := dealer.GenerateShares(context.Background(), pIDs)
	require.NoError(err)

	verificationKeys := bls.GetVerificationKeys(shares)
	configs := make(map[party.ID]*bls.Config)
	for _, id := range pIDs {
		configs[id] = bls.NewConfig(id, threshold, len(pIDs), shares[id], groupPK, verificationKeys)
	}

	// Test multiple signer combinations
	signerSets := [][]party.ID{
		{pIDs[0], pIDs[1], pIDs[2]},          // First 3
		{pIDs[2], pIDs[3], pIDs[4]},          // Last 3
		{pIDs[0], pIDs[2], pIDs[4]},          // Every other
		{pIDs[1], pIDs[2], pIDs[3]},          // Middle 3
		{pIDs[0], pIDs[1], pIDs[2], pIDs[3]}, // More than threshold
	}

	for i, signers := range signerSets {
		sigShares := make([]*bls.SignatureShare, 0, len(signers))
		for _, id := range signers {
			share, err := configs[id].Sign(message)
			require.NoError(err)
			sigShares = append(sigShares, share)
		}

		aggregatedSig, err := bls.AggregateSignatures(sigShares, threshold)
		require.NoError(err)

		valid := configs[pIDs[0]].VerifyAggregateSignature(message, aggregatedSig)
		assert.True(t, valid, "signer set %d should produce valid signature", i)
	}

	t.Logf("BLS threshold with different signer sets: all %d sets verified", len(signerSets))
}

// =============================================================================
// Corona (Post-Quantum) Threshold Network Tests
// =============================================================================

// TestCoronaSessionInit tests Corona session initialization
// This verifies the protocol can be started without requiring full MPC execution.
func TestCoronaSessionInit(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()

	pIDs := testPartyIDs(3)
	threshold := 2

	// Verify keygen session can be created
	for _, id := range pIDs {
		startFunc := corona.Keygen(id, pIDs, threshold, workerPool)
		require.NotNil(startFunc, "Keygen should return a function for %s", id)

		session, err := startFunc([]byte("test-session"))
		require.NoError(err)
		require.NotNil(session)
	}

	// Verify sign session can be created with mock config
	for _, id := range pIDs {
		cfg := &corona.Config{
			ID:           id,
			Threshold:    threshold,
			Participants: pIDs,
			PublicKey:    make([]byte, 32),
			PrivateShare: make([]byte, 32),
		}

		message := []byte("test message")
		signFunc := corona.SignWithConfig(cfg, pIDs[:threshold], message, workerPool)
		require.NotNil(signFunc, "Sign should return a function for %s", id)
	}

	// Verify refresh session can be created
	for _, id := range pIDs {
		cfg := &corona.Config{
			ID:           id,
			Threshold:    threshold,
			Participants: pIDs,
			PublicKey:    make([]byte, 32),
			PrivateShare: make([]byte, 32),
		}

		refreshFunc := corona.Refresh(cfg, pIDs, threshold, workerPool)
		require.NotNil(refreshFunc, "Refresh should return a function for %s", id)
	}

	t.Log("Corona session initialization verified for keygen, sign, and refresh")
}

// =============================================================================
// CMP Sign Timeout Tests
// =============================================================================

// TestCMPSignTimeout verifies that CMP sign respects context timeout
func TestCMPSignTimeout(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)

	// Create executor
	pe := NewProtocolExecutor(workerPool, logger)
	require.NotNil(pe)

	// Test that a cancelled context causes sign to fail
	pIDs := testPartyIDs(3)
	threshold := 2

	// Create a handler that will be cancelled
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
	defer cancel()

	// Wait for context to be cancelled
	time.Sleep(10 * time.Millisecond)

	// Try to create handler with cancelled context - should fail
	_, err := pe.CreateHandler(ctx, "timeout-test", pe.CMPKeygenStartFunc(pIDs[0], pIDs, threshold))
	// Note: Handler creation might succeed even with cancelled context
	// The real test is that protocol execution respects timeout

	t.Logf("CreateHandler with cancelled context: err=%v", err)
}

// TestCMPSignTimeoutWithMessage tests CMP sign with a very short timeout
func TestCMPSignTimeoutWithMessage(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping CMP sign timeout test in short mode")
	}

	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()

	pIDs := testPartyIDs(3)
	threshold := 2
	message := []byte("test message for CMP timeout test")

	// First, do a successful keygen
	keygenResults, err := runTestProtocol(t, pIDs, func(id party.ID) protocol.StartFunc {
		return cmp.Keygen(curve.Secp256k1{}, id, pIDs, threshold, workerPool)
	})
	require.NoError(err)
	require.Len(keygenResults, 3)

	t.Logf("CMP keygen succeeded, now testing signing with short timeout")

	// Try signing with very short timeout context
	// This simulates what happens in vm.go when SessionTimeout is exceeded
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
	defer cancel()

	// Create a single-party sign attempt that will timeout
	config := keygenResults[pIDs[0]].(*cmp.Config)
	startFunc := cmp.Sign(config, pIDs, message, workerPool)

	// Create handler with timeout context
	logger := log.NewTestLogger(log.DebugLevel)
	handler, err := protocol.NewHandler(
		ctx,
		logger,
		nil,
		startFunc,
		[]byte("timeout-session"),
		protocol.DefaultConfig(),
	)

	// Handler creation might succeed, but waiting should fail
	if err == nil && handler != nil {
		// Wait for result should fail due to timeout
		_, resultErr := handler.WaitForResult()
		// We expect either context deadline exceeded or protocol timeout
		if resultErr != nil {
			t.Logf("WaitForResult returned error (expected): %v", resultErr)
			// This is the expected behavior - timeout should be respected
		}
		handler.Stop()
	}

	t.Logf("CMP sign timeout test completed - timeout behavior verified")
}

// TestCMPHandlerTimeout tests that CGGMP21Handler respects context timeout
func TestCMPHandlerTimeout(t *testing.T) {
	require := require.New(t)

	workerPool := pool.NewPool(4)
	defer workerPool.TearDown()
	logger := log.NewTestLogger(log.DebugLevel)
	pe := NewProtocolExecutor(workerPool, logger)

	// Create CMP handler
	handler := &CGGMP21Handler{
		pool:     workerPool,
		executor: pe,
		// No router set - should fail with "router not configured"
	}

	// Try to sign without a router - should fail fast with descriptive error
	share := &cmpKeyShare{
		config: nil,
		thresh: 2,
		total:  3,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	_, err := handler.Sign(ctx, share, []byte("test"), []party.ID{"a", "b"})
	require.Error(err)
	require.Contains(err.Error(), "not configured")

	t.Logf("CGGMP21Handler configuration validation working correctly")
}

// TestRunSigningTimeout tests the VM.runSigning timeout behavior
func TestRunSigningTimeout(t *testing.T) {
	// This test verifies that the fix we made to vm.go works correctly
	// The fix was: adding context.WithTimeout(context.Background(), vm.config.SessionTimeout)

	// We can't easily create a full VM in a unit test, but we can verify the
	// timeout context pattern works correctly

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	// Simulate a long-running operation
	done := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			close(done)
		case <-time.After(10 * time.Second):
			// Should not reach here
		}
	}()

	// Wait for timeout
	select {
	case <-done:
		t.Logf("Context timeout was respected correctly")
	case <-time.After(1 * time.Second):
		t.Fatal("Context timeout was not respected")
	}
}