package ratelimit

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/time/rate"
)

func TestRateLimit(t *testing.T) {
	l := New(100, 1000)
	l.SetBackoff(time.Second, 10*time.Minute, 4, 0)

	now := time.Now()
	l.test(func() time.Time {
		return now
	})

	require.Equal(t, 999, l.N(999))
	require.Equal(t, 1, l.N(1000))

	// add 1 second to gain back limit
	now = now.Add(time.Second)
	require.Equal(t, 100, l.N(205))

	// add 2 seconds to take back limit * 2
	now = now.Add(2 * time.Second)
	require.Equal(t, 200, l.N(500))
	require.Equal(t, 0, l.N(1000))

	// allow enough time to burst and drain
	now = now.Add(10 * time.Second)
	require.Equal(t, 1000, l.N(1000))

	// allow enough time to burst and check amount after failure
	now = now.Add(10 * time.Second)
	l.Failure()
	require.Equal(t, 0, l.N(1000))

	// ensure burst had been temporarily set to the limit
	now = now.Add(4 * time.Second)
	require.Equal(t, 100, l.N(1000))

	// ensure burst is now fully available
	now = now.Add(10 * time.Second)
	require.Equal(t, 1000, l.N(1000))

	// reset failures
	l.Success()

	// allow enough time to burst and check amount after 10 failures
	now = now.Add(10 * time.Second)
	for i := 0; i < 10; i++ {
		l.Failure()
	}
	require.Equal(t, 0, l.N(1000))

	// delay for 10 minutes and check
	now = now.Add(10 * time.Minute)
	require.Equal(t, 0, l.N(1000))

	// reset failures
	now = now.Add(time.Second)
	l.Success()

	// after 1 second we should have more available
	now = now.Add(1 * time.Second)
	require.Equal(t, 100, l.N(1000))

	now = now.Add(5 * time.Second)
	require.Equal(t, 500, l.N(1000))
}

func TestExponentialBackoff(t *testing.T) {
	l := New(100, 1000)
	l.SetBackoff(time.Second, 10*time.Minute, 1.6, 0.2)

	now := time.Now()
	l.test(func() time.Time {
		return now
	})

	seconds := []int{1, 1, 2, 4, 6, 10, 16, 27, 43, 68, 110, 176, 281, 450, 600, 600, 600}
	for _, second := range seconds {
		l.Failure()

		l.Success() // should be ignored, as we're still waiting for the delay

		delay, newDelay := l.Delay()
		require.True(t, newDelay)
		require.Equal(t, time.Duration(second)*time.Second, delay.Truncate(time.Second), delay.Seconds())
	}

	// wait 10 minutes
	now = now.Add(10*time.Minute + time.Second)
	l.Success() // should now reset fails
	l.Failure()

	delay, newDelay := l.Delay()
	require.True(t, newDelay)
	require.Equal(t, time.Second, delay.Truncate(time.Second), delay.Seconds())
}

func TestSetLimitBurst(t *testing.T) {
	tests := map[string]struct {
		limit, burst  int
		expectedLimit rate.Limit
		expectedBurst int
	}{
		"unlimited limit (-1), unlimited burst (0)":              {-1, 0, rate.Inf, 0},
		"unlimited limit (0), unlimited burst (0)":               {0, 0, rate.Inf, 0},
		"unlimited limit (0), finite burst (100)":                {0, 100, rate.Inf, 100},
		"unlimited limit (-1), unlimited burst (-1)":             {-1, -1, rate.Inf, 0},
		"finite limit (100), zero burst (0) matches limit":       {100, 0, 100, 100},
		"finite limit (100), negative burst (-10) matches limit": {100, -10, 100, 100},
	}

	for tn, tc := range tests {
		t.Run(tn, func(t *testing.T) {
			l := New(tc.limit, tc.burst)
			assert.Equal(t, tc.expectedLimit, l.limiter.Limit())
			assert.Equal(t, tc.expectedBurst, l.limiter.Burst())
		})
	}
}
