package rate

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

	"github.com/clipperhouse/ntime"
	"github.com/stretchr/testify/require"
)

func TestLimiter_Wait(t *testing.T) {
	t.Parallel()

	t.Run("SingleBucket", func(t *testing.T) {
		t.Parallel()
		keyFunc := func(input string) string {
			return input
		}
		limit := NewLimit(2, 100*time.Millisecond)
		limiter := NewLimiter(keyFunc, limit)

		executionTime := ntime.Now()

		// Consume all tokens
		for range limit.Count() {
			ok := limiter.allow("test", executionTime)
			require.True(t, ok, "should allow initial tokens")
		}

		// Should not allow immediately
		ok := limiter.allow("test", executionTime)
		require.False(t, ok, "should not allow when tokens exhausted")

		// Test 1: Wait for 1 token with no cancellation
		{
			ctx := &testContext{
				done: make(chan struct{}), // never closes
			}

			allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, 1)
			require.True(t, allow, "should acquire 1 token after waiting")
			require.NotZero(t, details, "should return details")
			require.NoError(t, err, "should not return error")

			// We waited
			executionTime = executionTime.Add(limit.durationPerToken)
		}

		// Should not allow again immediately after wait
		ok = limiter.allow("test", executionTime)
		require.False(t, ok, "should not allow again immediately after wait")

		// Test 2: Wait for 1 token with immediate cancellation
		{
			// Context that's immediately cancelled (closed)
			done := make(chan struct{})
			ctx := &testContext{
				done: done,
			}
			close(done)

			allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, 1)
			require.False(t, allow, "should not acquire 1 token if context is cancelled immediately")
			require.NotZero(t, details, "should return details even on cancellation")
			require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
		}

		// Test 3: Wait for multiple tokens (limit.Count()) with no cancellation
		{
			// Wait for bucket to refill enough tokens
			executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.Count()))

			ctx := &testContext{
				done: make(chan struct{}), // never closes
			}

			allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, limit.Count())
			require.True(t, allow, "should acquire %d tokens after waiting", limit.Count())
			require.NotZero(t, details, "should return details")
			require.NoError(t, err, "should not return error")

			// We waited
			executionTime = executionTime.Add(limit.durationPerToken)
		}

		// Test 4: Wait for multiple tokens with immediate cancellation
		{
			// Wait for bucket to refill enough tokens
			executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.Count()))

			// Consume all tokens to exhaust the bucket again
			for range limit.Count() {
				ok := limiter.allow("test", executionTime)
				require.True(t, ok, "should allow tokens to exhaust bucket")
			}

			// Context that's immediately cancelled (closed)
			done := make(chan struct{})
			ctx := &testContext{
				done: done,
			}
			close(done)

			allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, limit.Count())
			require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.Count())
			require.NotZero(t, details, "should return details even on cancellation")
			require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
		}
	})

	t.Run("MultipleBuckets", func(t *testing.T) {
		t.Run("Serial", func(t *testing.T) {
			t.Parallel()
			keyFunc := func(input int) string {
				return fmt.Sprintf("test-bucket-%d", input)
			}
			const buckets = 3
			limit := NewLimit(2, 100*time.Millisecond)
			limiter := NewLimiter(keyFunc, limit)
			executionTime := ntime.Now()

			// Exhaust tokens for all buckets
			for bucketID := range buckets {
				for range limit.Count() {
					allow := limiter.allow(bucketID, executionTime)
					require.True(t, allow, "should allow initial tokens for bucket %d", bucketID)
				}
			}

			// Should not allow immediately for any bucket
			for bucketID := range buckets {
				allow := limiter.allow(bucketID, executionTime)
				require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID)
			}

			// Wait for a token for each bucket
			for bucketID := range buckets {
				ctx := &testContext{
					done: make(chan struct{}), // never closes
				}

				allow, details, err := limiter.waitNWithDetails(ctx, bucketID, executionTime, 1)
				require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID)
				require.NotZero(t, details, "should return details")
				require.NoError(t, err, "should not return error")
			}

			// We waited
			executionTime = executionTime.Add(limit.durationPerToken)

			// Buckets should be empty again
			for bucketID := range buckets {
				ok := limiter.allow(bucketID, executionTime)
				require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID)
			}
		})

		t.Run("Concurrent", func(t *testing.T) {
			t.Parallel()
			keyFunc := func(input int64) string {
				return fmt.Sprintf("test-bucket-%d", input)
			}
			const buckets int64 = 3
			limit := NewLimit(2, 100*time.Millisecond)
			limiter := NewLimiter(keyFunc, limit)
			executionTime := ntime.Now()

			// Exhaust tokens for all buckets
			for bucketID := range buckets {
				for range limit.Count() {
					allow := limiter.allow(bucketID, executionTime)
					require.True(t, allow, "should allow initial tokens for bucket %d", bucketID)
				}
			}

			// Buckets should be empty
			for bucketID := range buckets {
				allow := limiter.allow(bucketID, executionTime)
				require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID)
			}

			// Test 1: Multiple goroutines with immediate cancellation
			{
				concurrency := buckets
				results := make([]bool, concurrency)

				// Context that's immediately cancelled (closed)
				done := make(chan struct{})
				ctx := &testContext{
					done: done,
				}
				close(done)

				// Start concurrent waits
				var wg sync.WaitGroup
				for i := range concurrency {
					wg.Add(1)
					go func(i int64) {
						defer wg.Done()
						allow, _, err := limiter.waitNWithDetails(ctx, i, executionTime, 1)
						require.False(t, allow, "should not acquire token due to immediate cancellation")
						require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
						results[i] = allow
					}(i)
				}
				wg.Wait()

				// All should fail because context is cancelled immediately
				for i, result := range results {
					require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i)
				}
			}

			// Test 2: Multiple goroutines with delayed cancellation
			{
				keyFunc := func(input int64) int64 {
					return input
				}

				limit := NewLimit(1, 200*time.Millisecond)
				limiter := NewLimiter(keyFunc, limit)

				// Exhaust tokens
				for bucketID := range buckets {
					for range limit.Count() {
						allow := limiter.allow(bucketID, executionTime)
						require.True(t, allow, "should allow initial tokens for bucket %d", bucketID)
					}
				}

				// We're trying to emulate a wait call
				// that starts waiting with a retry, but
				// is cancelled in the meantime.

				const buckets int64 = 3
				results := make([]bool, buckets)

				var wg sync.WaitGroup
				for bucketID := range buckets {
					wg.Add(1)

					done := make(chan struct{})
					ctx := &testContext{
						done: done,
					}

					go func(bucketID int64) {
						defer wg.Done()
						// Should start waiting, since the bucket is empty.
						// retryAfter should be ~200ms
						allow, _, err := limiter.waitNWithDetails(ctx, bucketID, executionTime, 1)
						require.False(t, allow, "should not acquire token due to cancellation")
						require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
						results[bucketID] = allow
					}(bucketID)

					// Cancel context after goroutine has started.
					// Originally, this had a delay, for a better test,
					// but it is flaky on GitHub Actions, presumably
					// because the runner is resource constrained.
					// Delay works fine locally on Mac M2.
					close(done)
				}
				wg.Wait()

				// All should fail because context is cancelled
				for i, result := range results {
					require.False(t, result, "goroutine %d should not acquire token due to cancellation", i)
				}
			}
		})
	})
}

func TestLimiter_WaitWithDebug(t *testing.T) {
	t.Parallel()

	t.Run("SingleBucket", func(t *testing.T) {
		t.Parallel()
		keyFunc := func(input string) string {
			return input
		}
		limit := NewLimit(2, 100*time.Millisecond)
		limiter := NewLimiter(keyFunc, limit)

		executionTime := ntime.Now()

		// Consume all tokens
		for range limit.Count() {
			ok := limiter.allow("test", executionTime)
			require.True(t, ok, "should allow initial tokens")
		}

		// Should not allow immediately
		ok := limiter.allow("test", executionTime)
		require.False(t, ok, "should not allow when tokens exhausted")

		// Test 1: Wait for 1 token with no cancellation
		{
			ctx := &testContext{
				done: make(chan struct{}), // never closes
			}

			allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, 1)
			require.True(t, allow, "should acquire 1 token after waiting")
			require.NotNil(t, debugs, "should return debugs")
			require.NoError(t, err, "should not return error")

			// We waited
			executionTime = executionTime.Add(limit.durationPerToken)
		}

		// Should not allow again immediately after wait
		ok = limiter.allow("test", executionTime)
		require.False(t, ok, "should not allow again immediately after wait")

		// Test 2: Wait for 1 token with immediate cancellation
		{
			// Context that's immediately cancelled (closed)
			done := make(chan struct{})
			ctx := &testContext{
				done: done,
			}
			close(done)

			allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, 1)
			require.False(t, allow, "should not acquire 1 token if context is cancelled immediately")
			require.NotNil(t, debugs, "should return debugs even on cancellation")
			require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
		}

		// Test 3: Wait for multiple tokens (limit.Count()) with no cancellation
		{
			// Wait for bucket to refill enough tokens
			executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.Count()))

			ctx := &testContext{
				done: make(chan struct{}), // never closes
			}

			allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, limit.Count())
			require.True(t, allow, "should acquire %d tokens after waiting", limit.Count())
			require.NotNil(t, debugs, "should return debugs")
			require.NoError(t, err, "should not return error")

			// We waited
			executionTime = executionTime.Add(limit.durationPerToken)
		}

		// Test 4: Wait for multiple tokens with immediate cancellation
		{
			// Wait for bucket to refill enough tokens
			executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.Count()))

			// Consume all tokens to exhaust the bucket again
			for range limit.Count() {
				ok := limiter.allow("test", executionTime)
				require.True(t, ok, "should allow tokens to exhaust bucket")
			}

			// Context that's immediately cancelled (closed)
			done := make(chan struct{})
			ctx := &testContext{
				done: done,
			}
			close(done)

			allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, limit.Count())
			require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.Count())
			require.NotNil(t, debugs, "should return debugs even on cancellation")
			require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
		}
	})

	t.Run("MultipleBuckets", func(t *testing.T) {
		t.Run("Serial", func(t *testing.T) {
			t.Parallel()
			keyFunc := func(input int) string {
				return fmt.Sprintf("test-bucket-%d", input)
			}
			const buckets = 3
			limit := NewLimit(2, 100*time.Millisecond)
			limiter := NewLimiter(keyFunc, limit)
			executionTime := ntime.Now()

			// Exhaust tokens for all buckets
			for bucketID := range buckets {
				for range limit.Count() {
					allow := limiter.allow(bucketID, executionTime)
					require.True(t, allow, "should allow initial tokens for bucket %d", bucketID)
				}
			}

			// Should not allow immediately for any bucket
			for bucketID := range buckets {
				allow := limiter.allow(bucketID, executionTime)
				require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID)
			}

			// Wait for a token for each bucket
			for bucketID := range buckets {
				ctx := &testContext{
					done: make(chan struct{}), // never closes
				}

				allow, debugs, err := limiter.waitNWithDebug(ctx, bucketID, executionTime, 1)
				require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID)
				require.NotNil(t, debugs, "should return debugs")
				require.NoError(t, err, "should not return error")
			}

			// We waited
			executionTime = executionTime.Add(limit.durationPerToken)

			// Buckets should be empty again
			for bucketID := range buckets {
				ok := limiter.allow(bucketID, executionTime)
				require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID)
			}
		})

		t.Run("Concurrent", func(t *testing.T) {
			t.Parallel()
			keyFunc := func(input int64) string {
				return fmt.Sprintf("test-bucket-%d", input)
			}
			const buckets int64 = 3
			limit := NewLimit(2, 100*time.Millisecond)
			limiter := NewLimiter(keyFunc, limit)
			executionTime := ntime.Now()

			// Exhaust tokens for all buckets
			for bucketID := range buckets {
				for range limit.Count() {
					allow := limiter.allow(bucketID, executionTime)
					require.True(t, allow, "should allow initial tokens for bucket %d", bucketID)
				}
			}

			// Buckets should be empty
			for bucketID := range buckets {
				allow := limiter.allow(bucketID, executionTime)
				require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID)
			}

			// Test 1: Multiple goroutines with immediate cancellation
			{
				concurrency := buckets
				results := make([]bool, concurrency)

				// Context that's immediately cancelled (closed)
				done := make(chan struct{})
				ctx := &testContext{
					done: done,
				}
				close(done)

				// Start concurrent waits
				var wg sync.WaitGroup
				for i := range concurrency {
					wg.Add(1)
					go func(i int64) {
						defer wg.Done()
						allow, _, err := limiter.waitNWithDebug(ctx, i, executionTime, 1)
						require.False(t, allow, "should not acquire token due to immediate cancellation")
						require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
						results[i] = allow
					}(i)
				}
				wg.Wait()

				// All should fail because context is cancelled immediately
				for i, result := range results {
					require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i)
				}
			}

			// Test 2: Multiple goroutines with delayed cancellation
			{
				keyFunc := func(input int64) int64 {
					return input
				}

				limit := NewLimit(1, 200*time.Millisecond)
				limiter := NewLimiter(keyFunc, limit)

				// Exhaust tokens
				for bucketID := range buckets {
					for range limit.Count() {
						allow := limiter.allow(bucketID, executionTime)
						require.True(t, allow, "should allow initial tokens for bucket %d", bucketID)
					}
				}

				// We're trying to emulate a wait call
				// that starts waiting with a retry, but
				// is cancelled in the meantime.

				const buckets int64 = 3
				results := make([]bool, buckets)

				var wg sync.WaitGroup
				for bucketID := range buckets {
					wg.Add(1)

					done := make(chan struct{})
					ctx := &testContext{
						done: done,
					}

					go func(bucketID int64) {
						defer wg.Done()
						// Should start waiting, since the bucket is empty.
						// retryAfter should be ~200ms
						allow, _, err := limiter.waitNWithDebug(ctx, bucketID, executionTime, 1)
						require.False(t, allow, "should not acquire token due to cancellation")
						require.ErrorIs(t, err, context.Canceled, "should return context canceled error")
						results[bucketID] = allow
					}(bucketID)

					// Cancel context after goroutine has started.
					// Originally, this had a delay, for a better test,
					// but it is flaky on GitHub Actions, presumably
					// because the runner is resource constrained.
					// Delay works fine locally on Mac M2.
					close(done)
				}
				wg.Wait()

				// All should fail because context is cancelled
				for i, result := range results {
					require.False(t, result, "goroutine %d should not acquire token due to cancellation", i)
				}
			}
		})
	})
}

var _ context.Context = &testContext{}

type testContext struct {
	done <-chan struct{}
}

func (t *testContext) Deadline() (deadline time.Time, ok bool) {
	return time.Time{}, false
}

func (t *testContext) Done() <-chan struct{} {
	return t.done
}

func (t *testContext) Err() error {
	select {
	case <-t.done:
		return context.Canceled
	default:
		return nil
	}
}

func (t *testContext) Value(key any) any {
	return nil
}
