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

package wave

import (
	"context"
	"testing"
	"time"

	"github.com/luxfi/consensus/core/types"
	"github.com/luxfi/consensus/protocol/prism"
	"github.com/luxfi/consensus/protocol/wave/fpc"
	"github.com/stretchr/testify/require"
)

// mockTransport implements Transport for testing
type mockTransport[T comparable] struct {
	votes     map[T][]Photon[T]
	responses chan Photon[T]
}

func newMockTransport[T comparable]() *mockTransport[T] {
	return &mockTransport[T]{
		votes:     make(map[T][]Photon[T]),
		responses: make(chan Photon[T], 100),
	}
}

func (m *mockTransport[T]) RequestVotes(ctx context.Context, peers []types.NodeID, item T) <-chan Photon[T] {
	ch := make(chan Photon[T], len(m.votes[item]))
	go func() {
		defer close(ch)
		for _, vote := range m.votes[item] {
			select {
			case ch <- vote:
			case <-ctx.Done():
				return
			}
		}
	}()
	return ch
}

func (m *mockTransport[T]) MakeLocalPhoton(item T, prefer bool) Photon[T] {
	nodeID := [20]byte{1}
	return Photon[T]{
		Item:      item,
		Prefer:    prefer,
		Sender:    nodeID,
		Timestamp: time.Now(),
	}
}

func (m *mockTransport[T]) AddVote(item T, prefer bool) {
	if m.votes[item] == nil {
		m.votes[item] = make([]Photon[T], 0)
	}
	nodeID := [20]byte{byte(len(m.votes[item]) + 1)}
	m.votes[item] = append(m.votes[item], Photon[T]{
		Item:      item,
		Prefer:    prefer,
		Sender:    nodeID,
		Timestamp: time.Now(),
	})
}

// mockCut implements prism.Cut for testing
type mockCut[T comparable] struct {
	peers []types.NodeID
}

func newMockCut[T comparable](n int) *mockCut[T] {
	peers := make([]types.NodeID, n)
	for i := 0; i < n; i++ {
		nodeID := [20]byte{byte(i + 1)}
		peers[i] = nodeID
	}
	return &mockCut[T]{peers: peers}
}

func (m *mockCut[T]) Sample(k int) []types.NodeID {
	if k > len(m.peers) {
		k = len(m.peers)
	}
	return m.peers[:k]
}

func (m *mockCut[T]) Luminance() prism.Luminance {
	return prism.Luminance{
		ActivePeers: len(m.peers),
		TotalPeers:  len(m.peers),
		Lx:          float64(len(m.peers)),
	}
}

// TestWaveBasicVoting tests basic wave consensus voting
func TestWaveBasicVoting(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add unanimous votes for "tx1"
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
	}

	ctx := context.Background()
	wave.Tick(ctx, "tx1")

	// Check state after one round
	state, exists := wave.State("tx1")
	require.True(exists)
	require.False(state.Decided) // Need Beta rounds
	require.Equal(uint32(1), state.Count)
}

// TestWaveConfidenceBuilding tests confidence building over multiple rounds
func TestWaveConfidenceBuilding(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add unanimous votes
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
	}

	ctx := context.Background()

	// Run Beta rounds
	for i := uint32(0); i < cfg.Beta; i++ {
		wave.Tick(ctx, "tx1")
	}

	// Should be decided after Beta rounds
	state, exists := wave.State("tx1")
	require.True(exists)
	require.True(state.Decided)
	require.Equal(types.DecideAccept, state.Result)
}

// TestWavePreferenceSwitch tests preference switching
func TestWavePreferenceSwitch(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	ctx := context.Background()

	// First round: majority yes
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
	}
	wave.Tick(ctx, "tx1")

	state, _ := wave.State("tx1")
	require.Equal(uint32(1), state.Count)
	require.True(wave.Preference("tx1"))

	// Clear votes and add majority no
	tx.votes["tx1"] = nil
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", false)
	}
	wave.Tick(ctx, "tx1")

	// Count should reset on preference switch
	state, _ = wave.State("tx1")
	require.Equal(uint32(1), state.Count)
	require.False(wave.Preference("tx1"))
}

// TestWaveSplitVote tests handling of split votes
func TestWaveSplitVote(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       10,
		Alpha:   0.8, // Need 8 votes
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](15)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Split vote: 5 yes, 5 no (neither reaches threshold)
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
		tx.AddVote("tx1", false)
	}

	ctx := context.Background()
	wave.Tick(ctx, "tx1")

	// Count should reset on split vote
	state, _ := wave.State("tx1")
	require.Equal(uint32(0), state.Count)
}

// TestWaveMultipleItems tests handling multiple items concurrently
func TestWaveMultipleItems(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add votes for multiple items
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
		tx.AddVote("tx2", false)
		tx.AddVote("tx3", true)
	}

	ctx := context.Background()

	// Process all items
	wave.Tick(ctx, "tx1")
	wave.Tick(ctx, "tx2")
	wave.Tick(ctx, "tx3")

	// Check all states
	state1, _ := wave.State("tx1")
	state2, _ := wave.State("tx2")
	state3, _ := wave.State("tx3")

	require.Equal(uint32(1), state1.Count)
	require.True(wave.Preference("tx1"))

	require.Equal(uint32(1), state2.Count)
	require.False(wave.Preference("tx2"))

	require.Equal(uint32(1), state3.Count)
	require.True(wave.Preference("tx3"))
}

// TestWaveTimeout tests handling of vote timeout
func TestWaveTimeout(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 50 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Don't add any votes (will timeout)
	ctx := context.Background()
	wave.Tick(ctx, "tx1")

	// Should have state but no progress
	_, exists := wave.State("tx1")
	require.True(exists)
}

// TestWaveDecisionPersistence tests that decisions persist
func TestWaveDecisionPersistence(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add unanimous votes
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
	}

	ctx := context.Background()

	// Reach decision
	for i := uint32(0); i < cfg.Beta; i++ {
		wave.Tick(ctx, "tx1")
	}

	state, _ := wave.State("tx1")
	require.True(state.Decided)

	// Clear votes and add opposite votes
	tx.votes["tx1"] = nil
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", false)
	}

	// Additional tick should not change decision
	wave.Tick(ctx, "tx1")

	state, _ = wave.State("tx1")
	require.True(state.Decided)
	require.Equal(types.DecideAccept, state.Result) // Original decision persists
}

// TestWaveHighThroughput tests wave under high transaction load
func TestWaveHighThroughput(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](100)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	ctx := context.Background()

	// Process many items
	for i := 0; i < 100; i++ {
		item := string(rune('a' + i))
		for j := 0; j < 5; j++ {
			tx.AddVote(item, true)
		}
		wave.Tick(ctx, item)
	}

	// Verify all processed
	for i := 0; i < 100; i++ {
		item := string(rune('a' + i))
		_, exists := wave.State(item)
		require.True(exists)
	}
}

// TestWaveFPCWithCustomThresholds tests wave consensus with FPC enabled and custom thresholds
func TestWaveFPCWithCustomThresholds(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:         5,
		Alpha:     0.8, // Ignored when FPC is enabled
		Beta:      3,
		RoundTO:   100 * time.Millisecond,
		EnableFPC: true,
		ThetaMin:  0.6,
		ThetaMax:  0.9,
		FPCSeed:   fpc.DeriveEpochSeed(1, []byte("test"), nil),
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add unanimous votes
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
	}

	ctx := context.Background()

	// Run Beta rounds
	for i := uint32(0); i < cfg.Beta; i++ {
		wave.Tick(ctx, "tx1")
	}

	// Should reach decision with FPC
	state, exists := wave.State("tx1")
	require.True(exists)
	require.True(state.Decided)
	require.Equal(types.DecideAccept, state.Result)
}

// TestWaveFPCZeroThresholds tests FPC with zero thresholds (defaults applied)
func TestWaveFPCZeroThresholds(t *testing.T) {
	require := require.New(t)

	// ThetaMin and ThetaMax are 0, so defaults should be used
	cfg := Config{
		K:         5,
		Alpha:     0.8,
		Beta:      3,
		RoundTO:   100 * time.Millisecond,
		EnableFPC: true,
		ThetaMin:  0, // Should default to 0.5
		ThetaMax:  0, // Should default to 0.8
		FPCSeed:   fpc.DeriveEpochSeed(1, []byte("test"), nil),
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add unanimous votes
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", true)
	}

	ctx := context.Background()
	wave.Tick(ctx, "tx1")

	// Check state after one round
	state, exists := wave.State("tx1")
	require.True(exists)
	require.Equal(uint32(1), state.Count)
}

// TestWaveContextCancellation tests handling of context cancellation
func TestWaveContextCancellation(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 5 * time.Second, // Very long timeout (never reached)
	}

	cut := newMockCut[string](10)
	// Use slow transport that never sends votes
	tx := newSlowMockTransport[string](10 * time.Second)

	wave, _ := New[string](cfg, cut, tx)

	// Cancel context quickly to trigger ctx.Done() path
	ctx, cancel := context.WithCancel(context.Background())

	// Cancel context very quickly
	go func() {
		time.Sleep(5 * time.Millisecond)
		cancel()
	}()

	wave.Tick(ctx, "tx1")

	// State should exist but not be decided (early return from ctx.Done)
	state, exists := wave.State("tx1")
	require.True(exists)
	require.False(state.Decided)
}

// TestWaveRejectDecision tests reaching a reject decision
func TestWaveRejectDecision(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Add unanimous NO votes
	for i := 0; i < 5; i++ {
		tx.AddVote("tx1", false)
	}

	ctx := context.Background()

	// Run Beta rounds
	for i := uint32(0); i < cfg.Beta; i++ {
		wave.Tick(ctx, "tx1")
	}

	// Should reach reject decision
	state, exists := wave.State("tx1")
	require.True(exists)
	require.True(state.Decided)
	require.Equal(types.DecideReject, state.Result)
	require.False(wave.Preference("tx1"))
}

// TestWaveNoVotesTimeout tests handling when channel closes with no votes
func TestWaveNoVotesTimeout(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 10 * time.Millisecond, // Very short timeout
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// No votes added - the mock returns an empty closed channel
	// which times out, resulting in totalVotes=0 and early return
	ctx := context.Background()
	wave.Tick(ctx, "tx1")

	// State should exist
	state, exists := wave.State("tx1")
	require.True(exists)
	require.False(state.Decided)
}

// slowMockTransport delays sending votes to trigger timeout
type slowMockTransport[T comparable] struct {
	delay time.Duration
	votes []Photon[T]
}

func newSlowMockTransport[T comparable](delay time.Duration) *slowMockTransport[T] {
	return &slowMockTransport[T]{
		delay: delay,
		votes: make([]Photon[T], 0),
	}
}

func (m *slowMockTransport[T]) RequestVotes(ctx context.Context, peers []types.NodeID, item T) <-chan Photon[T] {
	ch := make(chan Photon[T], len(m.votes))
	go func() {
		defer close(ch)
		// Delay before sending any votes - may timeout first
		time.Sleep(m.delay)
		for _, vote := range m.votes {
			select {
			case ch <- vote:
			case <-ctx.Done():
				return
			}
		}
	}()
	return ch
}

func (m *slowMockTransport[T]) MakeLocalPhoton(item T, prefer bool) Photon[T] {
	nodeID := [20]byte{1}
	return Photon[T]{
		Item:      item,
		Prefer:    prefer,
		Sender:    nodeID,
		Timestamp: time.Now(),
	}
}

// TestWaveTimeoutPath tests the timeout path in vote collection
func TestWaveTimeoutPath(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 10 * time.Millisecond, // Short timeout
	}

	cut := newMockCut[string](10)
	// Use slow transport that delays 50ms (longer than 10ms timeout)
	tx := newSlowMockTransport[string](50 * time.Millisecond)

	wave, _ := New[string](cfg, cut, tx)

	ctx := context.Background()
	wave.Tick(ctx, "tx1")

	// State should exist but no progress (timeout with 0 votes)
	state, exists := wave.State("tx1")
	require.True(exists)
	require.False(state.Decided)
}

// TestWaveStateNotExists tests State for non-existent item
func TestWaveStateNotExists(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Query state for item that was never processed
	state, exists := wave.State("nonexistent")
	require.False(exists)
	require.Nil(state)
}

// TestWavePreferenceNotExists tests Preference for non-existent item
func TestWavePreferenceNotExists(t *testing.T) {
	require := require.New(t)

	cfg := Config{
		K:       5,
		Alpha:   0.8,
		Beta:    3,
		RoundTO: 100 * time.Millisecond,
	}

	cut := newMockCut[string](10)
	tx := newMockTransport[string]()
	wave, _ := New[string](cfg, cut, tx)

	// Query preference for item that was never processed
	pref := wave.Preference("nonexistent")
	require.False(pref) // Default value for bool
}