package limithandler

import (
	"context"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
	"gitlab.com/gitlab-org/gitaly/v14/internal/metadata/featureflag"
	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
)

type counter struct {
	sync.Mutex
	max         int
	current     int
	queued      int
	dequeued    int
	enter       int
	exit        int
	droppedSize int
	droppedTime int
}

func (c *counter) up() {
	c.Lock()
	defer c.Unlock()

	c.current = c.current + 1
	if c.current > c.max {
		c.max = c.current
	}
}

func (c *counter) down() {
	c.Lock()
	defer c.Unlock()

	c.current = c.current - 1
}

func (c *counter) currentVal() int {
	c.Lock()
	defer c.Unlock()
	return c.current
}

func (c *counter) Queued(ctx context.Context) {
	c.Lock()
	defer c.Unlock()
	c.queued++
}

func (c *counter) Dequeued(ctx context.Context) {
	c.Lock()
	defer c.Unlock()
	c.dequeued++
}

func (c *counter) Enter(ctx context.Context, acquireTime time.Duration) {
	c.Lock()
	defer c.Unlock()
	c.enter++
}

func (c *counter) Exit(ctx context.Context) {
	c.Lock()
	defer c.Unlock()
	c.exit++
}

func (c *counter) Dropped(ctx context.Context, reason string) {
	switch reason {
	case "max_time":
		c.droppedTime++
	case "max_size":
		c.droppedSize++
	}
}

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

	tests := []struct {
		name             string
		concurrency      int
		maxConcurrency   int
		iterations       int
		buckets          int
		wantMonitorCalls bool
	}{
		{
			name:             "single",
			concurrency:      1,
			maxConcurrency:   1,
			iterations:       1,
			buckets:          1,
			wantMonitorCalls: true,
		},
		{
			name:             "two-at-a-time",
			concurrency:      100,
			maxConcurrency:   2,
			iterations:       10,
			buckets:          1,
			wantMonitorCalls: true,
		},
		{
			name:             "two-by-two",
			concurrency:      100,
			maxConcurrency:   2,
			iterations:       4,
			buckets:          2,
			wantMonitorCalls: true,
		},
		{
			name:             "no-limit",
			concurrency:      10,
			maxConcurrency:   0,
			iterations:       200,
			buckets:          1,
			wantMonitorCalls: false,
		},
		{
			name:           "wide-spread",
			concurrency:    1000,
			maxConcurrency: 2,
			// We use a long delay here to prevent flakiness in CI. If the delay is
			// too short, the first goroutines to enter the critical section will be
			// gone before we hit the intended maximum concurrency.
			iterations:       40,
			buckets:          50,
			wantMonitorCalls: true,
		},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			ctx := testhelper.Context(t)

			expectedGaugeMax := tt.maxConcurrency * tt.buckets
			if tt.maxConcurrency <= 0 {
				expectedGaugeMax = tt.concurrency
			}

			gauge := &counter{}

			limiter := NewLimiter(
				tt.maxConcurrency,
				0,
				nil,
				gauge,
			)
			wg := sync.WaitGroup{}
			wg.Add(tt.concurrency)

			full := sync.NewCond(&sync.Mutex{})

			// primePump waits for the gauge to reach the minimum
			// expected max concurrency so that the limiter is
			// "warmed" up before proceeding with the test
			primePump := func() {
				full.L.Lock()
				defer full.L.Unlock()

				gauge.up()

				if gauge.max >= expectedGaugeMax {
					full.Broadcast()
					return
				}

				full.Wait() // wait until full is broadcast
			}

			// We know of an edge case that can lead to the rate limiter
			// occasionally letting one or two extra goroutines run
			// concurrently.
			for c := 0; c < tt.concurrency; c++ {
				go func(counter int) {
					for i := 0; i < tt.iterations; i++ {
						lockKey := strconv.Itoa((i ^ counter) % tt.buckets)

						_, err := limiter.Limit(ctx, lockKey, func() (interface{}, error) {
							primePump()

							current := gauge.currentVal()
							require.True(t, current <= expectedGaugeMax, "Expected the number of concurrent operations (%v) to not exceed the maximum concurrency (%v)", current, expectedGaugeMax)

							require.True(t, limiter.countSemaphores() <= tt.buckets, "Expected the number of semaphores (%v) to be lte number of buckets (%v)", limiter.countSemaphores(), tt.buckets)

							gauge.down()
							return nil, nil
						})
						require.NoError(t, err)
					}

					wg.Done()
				}(c)
			}

			wg.Wait()

			assert.Equal(t, expectedGaugeMax, gauge.max, "Expected maximum concurrency")
			assert.Equal(t, 0, gauge.current)
			assert.Equal(t, 0, limiter.countSemaphores())

			var wantMonitorCallCount int
			if tt.wantMonitorCalls {
				wantMonitorCallCount = tt.concurrency * tt.iterations
			} else {
				wantMonitorCallCount = 0
			}

			assert.Equal(t, wantMonitorCallCount, gauge.enter)
			assert.Equal(t, wantMonitorCallCount, gauge.exit)
			assert.Equal(t, wantMonitorCallCount, gauge.queued)
			assert.Equal(t, wantMonitorCallCount, gauge.dequeued)
		})
	}
}

type blockingQueueCounter struct {
	counter

	queuedCh chan struct{}
}

// Queued will block on a channel. We need a way to synchronize on when a Limiter has attempted to acquire
// a semaphore but has not yet. The caller can use the channel to wait for many requests to be queued
func (b *blockingQueueCounter) Queued(_ context.Context) {
	b.queuedCh <- struct{}{}
}

func TestConcurrencyLimiter_queueLimit(t *testing.T) {
	queueLimit := 10
	ctx := testhelper.Context(t)

	testCases := []struct {
		desc          string
		featureFlagOn bool
	}{
		{
			desc:          "feature flag on",
			featureFlagOn: true,
		},
		{
			desc:          "feature flag off",
			featureFlagOn: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			ctx = featureflag.IncomingCtxWithFeatureFlag(
				ctx,
				featureflag.ConcurrencyQueueEnforceMax,
				tc.featureFlagOn,
			)

			monitorCh := make(chan struct{})
			monitor := &blockingQueueCounter{queuedCh: monitorCh}
			ch := make(chan struct{})
			limiter := NewLimiter(1, queueLimit, nil, monitor)

			// occupied with one live request that takes a long time to complete
			go func() {
				_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
					ch <- struct{}{}
					<-ch
					return nil, nil
				})
				require.NoError(t, err)
			}()

			<-monitorCh
			<-ch

			var wg sync.WaitGroup
			// fill up the queue
			for i := 0; i < queueLimit; i++ {
				wg.Add(1)
				go func() {
					_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
						return nil, nil
					})
					require.NoError(t, err)
					wg.Done()
				}()
			}

			var queued int
			for range monitorCh {
				queued++
				if queued == queueLimit {
					break
				}
			}

			errChan := make(chan error, 1)
			go func() {
				_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
					return nil, nil
				})
				errChan <- err
			}()

			if tc.featureFlagOn {
				err := <-errChan
				assert.Error(t, err)
				assert.Equal(t, "maximum queue size reached", err.Error())
				assert.Equal(t, monitor.droppedSize, 1)
			} else {
				<-monitorCh
				assert.Equal(t, int64(queueLimit+1), limiter.queued)
				assert.Equal(t, monitor.droppedSize, 0)
			}

			close(ch)
			wg.Wait()
		})
	}
}

type blockingDequeueCounter struct {
	counter

	dequeuedCh chan struct{}
}

// Dequeued will block on a channel. We need a way to synchronize on when a Limiter has successfully
// acquired a semaphore but has not yet. The caller can use the channel to wait for many requests to
// be queued
func (b *blockingDequeueCounter) Dequeued(context.Context) {
	b.dequeuedCh <- struct{}{}
}

func TestLimitConcurrency_queueWaitTime(t *testing.T) {
	ctx := testhelper.Context(t)

	t.Run("with feature flag", func(t *testing.T) {
		ctx = featureflag.IncomingCtxWithFeatureFlag(
			ctx,
			featureflag.ConcurrencyQueueMaxWait,
			true,
		)

		ticker := helper.NewManualTicker()

		dequeuedCh := make(chan struct{})
		monitor := &blockingDequeueCounter{dequeuedCh: dequeuedCh}

		limiter := NewLimiter(
			1,
			0,
			func() helper.Ticker {
				return ticker
			},
			monitor,
		)

		ch := make(chan struct{})
		var wg sync.WaitGroup
		wg.Add(1)
		go func() {
			_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
				<-ch
				return nil, nil
			})
			require.NoError(t, err)
			wg.Done()
		}()

		<-dequeuedCh

		ticker.Tick()

		errChan := make(chan error)
		go func() {
			_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
				return nil, nil
			})
			errChan <- err
		}()

		<-dequeuedCh
		err := <-errChan

		assert.Equal(t, ErrMaxQueueTime, err)
		assert.Equal(t, monitor.droppedTime, 1)
		close(ch)
		wg.Wait()
	})

	t.Run("without feature flag", func(t *testing.T) {
		ctx = featureflag.IncomingCtxWithFeatureFlag(
			ctx,
			featureflag.ConcurrencyQueueMaxWait,
			false,
		)

		ticker := helper.NewManualTicker()

		dequeuedCh := make(chan struct{})
		monitor := &blockingDequeueCounter{dequeuedCh: dequeuedCh}

		limiter := NewLimiter(
			1,
			0,
			func() helper.Ticker {
				return ticker
			},
			monitor,
		)

		ch := make(chan struct{})
		go func() {
			_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
				<-ch
				return nil, nil
			})
			require.NoError(t, err)
		}()

		<-dequeuedCh

		ticker.Tick()

		errChan := make(chan error)
		go func() {
			_, err := limiter.Limit(ctx, "key", func() (interface{}, error) {
				return nil, nil
			})
			errChan <- err
		}()

		close(ch)
		<-dequeuedCh
		err := <-errChan

		assert.NoError(t, err)
		assert.Equal(t, monitor.droppedTime, 0)
	})
}
