//go:build go1.25

package messagequeue

import (
	"context"
	"errors"
	"sync"
	"testing"
	"testing/synctest"
	"time"

	cid "github.com/ipfs/go-cid"
	"github.com/ipfs/go-test/random"
	"github.com/libp2p/go-libp2p/p2p/protocol/ping"
)

type mockPeerConn struct {
	err       error
	latency   time.Duration
	latencies []time.Duration
	pinged    chan struct{}
}

func (pc *mockPeerConn) Ping(ctx context.Context) ping.Result {
	timer := time.NewTimer(pc.latency)
	pc.pinged <- struct{}{}
	select {
	case <-timer.C:
		if pc.err != nil {
			return ping.Result{Error: pc.err}
		}
		pc.latencies = append(pc.latencies, pc.latency)
	case <-ctx.Done():
	}
	return ping.Result{RTT: pc.latency}
}

func (pc *mockPeerConn) Latency() time.Duration {
	sum := time.Duration(0)
	if len(pc.latencies) == 0 {
		return sum
	}
	for _, l := range pc.latencies {
		sum += l
	}
	return sum / time.Duration(len(pc.latencies))
}

type timeoutRecorder struct {
	timedOutKs []cid.Cid
	lk         sync.Mutex
}

func (tr *timeoutRecorder) onTimeout(tks []cid.Cid, _ time.Duration) {
	tr.lk.Lock()
	defer tr.lk.Unlock()

	tr.timedOutKs = append(tr.timedOutKs, tks...)
}

func (tr *timeoutRecorder) timedOutCount() int {
	tr.lk.Lock()
	defer tr.lk.Unlock()

	return len(tr.timedOutKs)
}

func (tr *timeoutRecorder) clear() {
	tr.lk.Lock()
	defer tr.lk.Unlock()

	tr.timedOutKs = nil
}

func TestDontHaveTimeoutMgrTimeout(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		firstks := random.Cids(2)
		secondks := append(firstks, random.Cids(3)...)
		latency := time.Millisecond * 20
		latMultiplier := 2
		expProcessTime := 5 * time.Millisecond
		expectedTimeout := expProcessTime + latency*time.Duration(latMultiplier)
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})
		cfg := DefaultDontHaveTimeoutConfig()
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond
		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged
		// Add first set of keys
		dhtm.AddPending(firstks)

		// Wait for less than the expected timeout
		time.Sleep(expectedTimeout - 10*time.Millisecond)

		// At this stage no keys should have timed out
		if tr.timedOutCount() > 0 {
			t.Fatal("expected timeout not to have happened yet")
		}

		// Add second set of keys
		dhtm.AddPending(secondks)

		// Wait until after the expected timeout
		time.Sleep(20 * time.Millisecond)
		synctest.Wait()

		canRetry := true

	retry:
		select {
		case <-timeoutsTriggered:
		case <-time.After(2 * time.Second):
			if canRetry {
				canRetry = false
				time.Sleep(10 * time.Millisecond)
				goto retry
			}
			t.Fatal("timed out waiting for timeouts")
		}

		// At this stage first set of keys should have timed out
		if tr.timedOutCount() != len(firstks) {
			t.Fatal("expected timeout", tr.timedOutCount(), len(firstks))
		}
		// Clear the recorded timed out keys
		tr.clear()

		// Sleep until the second set of keys should have timed out
		time.Sleep(expectedTimeout + 10*time.Millisecond)

		<-timeoutsTriggered

		// At this stage all keys should have timed out. The second set included
		// the first set of keys, but they were added before the first set timed
		// out, so only the remaining keys should have beed added.
		if tr.timedOutCount() != len(secondks)-len(firstks) {
			t.Fatal("expected second set of keys to timeout")
		}
	})
}

func TestDontHaveTimeoutMgrCancel(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(3)
		latency := time.Millisecond * 10
		latMultiplier := 1
		expProcessTime := time.Duration(0)
		expectedTimeout := latency
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})
		cfg := DefaultDontHaveTimeoutConfig()
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond
		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged

		// Add keys
		dhtm.AddPending(ks)
		time.Sleep(5 * time.Millisecond)
		synctest.Wait()

		// Cancel keys
		cancelCount := 1
		dhtm.CancelPending(ks[:cancelCount])

		// Wait for the expected timeout
		time.Sleep(expectedTimeout)

		<-timeoutsTriggered

		// At this stage all non-cancelled keys should have timed out
		if tr.timedOutCount() != len(ks)-cancelCount {
			t.Fatal("expected timeout")
		}
	})
}

func TestDontHaveTimeoutWantCancelWant(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(3)
		latency := time.Millisecond * 20
		latMultiplier := 1
		expProcessTime := time.Duration(0)
		expectedTimeout := latency
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})
		cfg := DefaultDontHaveTimeoutConfig()
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond
		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged

		// Add keys
		dhtm.AddPending(ks)

		// Wait for a short time
		time.Sleep(expectedTimeout - 10*time.Millisecond)
		synctest.Wait()

		// Cancel two keys
		dhtm.CancelPending(ks[:2])

		time.Sleep(5 * time.Millisecond)
		synctest.Wait()

		// Add back one cancelled key
		dhtm.AddPending(ks[:1])

		// Wait till after initial timeout
		time.Sleep(10 * time.Millisecond)

		<-timeoutsTriggered

		// At this stage only the key that was never cancelled should have timed out
		if tr.timedOutCount() != 1 {
			t.Fatal("expected one key to timeout")
		}

		// Wait till after added back key should time out
		time.Sleep(latency)

		<-timeoutsTriggered

		// At this stage the key that was added back should also have timed out
		if tr.timedOutCount() != 2 {
			t.Fatal("expected added back key to timeout")
		}
	})
}

func TestDontHaveTimeoutRepeatedAddPending(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(10)
		latency := time.Millisecond * 5
		latMultiplier := 1
		expProcessTime := time.Duration(0)
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})

		cfg := DefaultDontHaveTimeoutConfig()
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond

		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged

		// Add keys repeatedly
		for _, c := range ks {
			dhtm.AddPending([]cid.Cid{c})
		}

		// Wait for the expected timeout
		time.Sleep(latency + 5*time.Millisecond)

		<-timeoutsTriggered

		// At this stage all keys should have timed out
		if tr.timedOutCount() != len(ks) {
			t.Fatal("expected timeout")
		}
	})
}

func TestDontHaveTimeoutMgrMessageLatency(t *testing.T) {
	t.Parallel()
	ks := random.Cids(2)
	synctest.Test(t, func(t *testing.T) {
		latency := time.Millisecond * 40
		latMultiplier := 1
		expProcessTime := time.Duration(0)
		msgLatencyMultiplier := 1
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})

		cfg := DefaultDontHaveTimeoutConfig()
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MessageLatencyMultiplier = msgLatencyMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond

		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged
		// Add keys
		dhtm.AddPending(ks)

		// expectedTimeout
		// = expProcessTime + latency*time.Duration(latMultiplier)
		// = 0 + 40ms * 1
		// = 40ms

		// Wait for less than the expected timeout
		time.Sleep(25 * time.Millisecond)
		synctest.Wait()

		// Receive two message latency updates
		dhtm.UpdateMessageLatency(time.Millisecond * 20)
		dhtm.UpdateMessageLatency(time.Millisecond * 10)

		// alpha is 0.5 so timeout should be
		// = (20ms * alpha) + (10ms * (1 - alpha))
		// = (20ms * 0.5) + (10ms * 0.5)
		// = 15ms
		// We've already slept for 25ms so with the new 15ms timeout
		// the keys should have timed out

		// Give the queue some time to process the updates
		time.Sleep(5 * time.Millisecond)

		<-timeoutsTriggered

		if tr.timedOutCount() != len(ks) {
			t.Fatal("expected keys to timeout")
		}
	})
}

func TestDontHaveTimeoutMgrMessageLatencyMax(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(2)
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: time.Second,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		msgLatencyMultiplier := 1
		testMaxTimeout := time.Millisecond * 10
		timeoutsTriggered := make(chan struct{})

		cfg := DefaultDontHaveTimeoutConfig()
		cfg.MessageLatencyMultiplier = msgLatencyMultiplier
		cfg.MaxTimeout = testMaxTimeout
		cfg.MinTimeout = 10 * time.Millisecond

		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged
		// Add keys
		dhtm.AddPending(ks)

		// Receive a message latency update that would make the timeout greater
		// than the maximum timeout
		dhtm.UpdateMessageLatency(testMaxTimeout * 4)

		// Sleep until just after the maximum timeout
		time.Sleep(testMaxTimeout + 5*time.Millisecond)

		<-timeoutsTriggered

		// Keys should have timed out
		if tr.timedOutCount() != len(ks) {
			t.Fatal("expected keys to timeout")
		}
	})
}

func TestDontHaveTimeoutMgrUsesDefaultTimeoutIfPingError(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(2)
		latency := time.Millisecond
		latMultiplier := 2
		expProcessTime := 2 * time.Millisecond
		defaultTimeout := 10 * time.Millisecond
		expectedTimeout := expProcessTime + defaultTimeout
		tr := timeoutRecorder{}
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
			err:     errors.New("ping error"),
		}
		timeoutsTriggered := make(chan struct{})

		cfg := DefaultDontHaveTimeoutConfig()
		cfg.DontHaveTimeout = defaultTimeout
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond
		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged

		// Add keys
		dhtm.AddPending(ks)

		// Sleep for less than the expected timeout
		time.Sleep(expectedTimeout - 5*time.Millisecond)
		synctest.Wait()

		// At this stage no timeout should have happened yet
		if tr.timedOutCount() > 0 {
			t.Fatal("expected timeout not to have happened yet")
		}

		// Sleep until after the expected timeout
		time.Sleep(10 * time.Millisecond)

		<-timeoutsTriggered

		// Now the keys should have timed out
		if tr.timedOutCount() != len(ks) {
			t.Fatal("expected timeout")
		}
	})
}

func TestDontHaveTimeoutMgrUsesDefaultTimeoutIfLatencyLonger(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(2)
		latency := time.Millisecond * 200
		latMultiplier := 1
		expProcessTime := time.Duration(0)
		defaultTimeout := 100 * time.Millisecond
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})

		cfg := DefaultDontHaveTimeoutConfig()
		cfg.DontHaveTimeout = defaultTimeout
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond

		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged

		// Add keys
		dhtm.AddPending(ks)

		// Sleep for less than the default timeout
		time.Sleep(defaultTimeout - 50*time.Millisecond)
		synctest.Wait()

		// At this stage no timeout should have happened yet
		if tr.timedOutCount() > 0 {
			t.Fatal("expected timeout not to have happened yet")
		}

		// Sleep until after the default timeout
		time.Sleep(defaultTimeout * 2)

		<-timeoutsTriggered

		// Now the keys should have timed out
		if tr.timedOutCount() != len(ks) {
			t.Fatal("expected timeout")
		}
	})
}

func TestDontHaveTimeoutNoTimeoutAfterShutdown(t *testing.T) {
	t.Parallel()
	synctest.Test(t, func(t *testing.T) {
		ks := random.Cids(2)
		latency := time.Millisecond * 10
		latMultiplier := 1
		expProcessTime := time.Duration(0)
		pinged := make(chan struct{})
		pc := &mockPeerConn{
			latency: latency,
			pinged:  pinged,
		}
		tr := timeoutRecorder{}
		timeoutsTriggered := make(chan struct{})

		cfg := DefaultDontHaveTimeoutConfig()
		cfg.PingLatencyMultiplier = latMultiplier
		cfg.MaxExpectedWantProcessTime = expProcessTime
		cfg.MinTimeout = 10 * time.Millisecond
		cfg.timeoutsSignal = timeoutsTriggered
		dhtm := newDontHaveTimeoutMgr(pc, tr.onTimeout, cfg)
		dhtm.Start()
		defer dhtm.Shutdown()
		<-pinged

		// Add keys
		dhtm.AddPending(ks)

		// Wait less than the timeout
		time.Sleep(latency - 5*time.Millisecond)
		synctest.Wait()

		// Shutdown the manager
		dhtm.Shutdown()

		// Wait for the expected timeout
		time.Sleep(10 * time.Millisecond)
		synctest.Wait()

		// Manager was shut down so timeout should not have fired
		if tr.timedOutCount() != 0 {
			t.Fatal("expected no timeout after shutdown")
		}
	})
}
