package limiting_test

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting"
	"github.com/stretchr/testify/require"
)

const (
	operationDuration = 100 * time.Millisecond
	operationCount    = 64
)

type testKeyLimit struct {
	keys         []string
	limit        int
	withoutLimit bool
	failCount    atomic.Int32
}

func TestLimiter(t *testing.T) {
	t.Run("duplicate key", func(t *testing.T) {
		_, err := limiting.New([]limiting.KeyLimit{
			{[]string{"A", "B"}, 10},
			{[]string{"B", "C"}, 10},
		})
		require.Error(t, err)
	})

	testLimits := []*testKeyLimit{
		{keys: []string{"A"}, limit: operationCount / 4},
		{keys: []string{"B"}, limit: operationCount / 2},
		{keys: []string{"C", "D"}, limit: operationCount / 4},
		{keys: []string{"E"}, limit: 2 * operationCount},
		{keys: []string{"F"}, withoutLimit: true},
	}

	t.Run("non-blocking mode", func(t *testing.T) {
		testLimiter(t, testLimits, false)
	})

	resetFailCounts(testLimits)

	t.Run("blocking mode", func(t *testing.T) {
		testLimiter(t, testLimits, true)
	})
}

func testLimiter(t *testing.T, testCases []*testKeyLimit, blocking bool) {
	lr, err := limiting.New(getLimits(testCases))
	require.NoError(t, err)

	tasks := createTestTasks(testCases, lr, blocking)

	t.Run("first run", func(t *testing.T) {
		executeTasks(tasks...)
		verifyResults(t, testCases, blocking)
	})

	t.Run("repeated run", func(t *testing.T) {
		resetFailCounts(testCases)
		executeTasks(tasks...)
		verifyResults(t, testCases, blocking)
	})
}

func getLimits(testCases []*testKeyLimit) []limiting.KeyLimit {
	var limits []limiting.KeyLimit
	for _, tc := range testCases {
		if tc.withoutLimit {
			continue
		}
		limits = append(limits, limiting.KeyLimit{
			Keys:  tc.keys,
			Limit: uint64(tc.limit),
		})
	}
	return limits
}

func createTestTasks(testCases []*testKeyLimit, lr *limiting.Limiter, blocking bool) []func() {
	var tasks []func()
	for _, tc := range testCases {
		for _, key := range tc.keys {
			tasks = append(tasks, func() {
				executeTaskN(operationCount, func() { acquireAndExecute(tc, lr, key, blocking) })
			})
		}
	}
	return tasks
}

func acquireAndExecute(tc *testKeyLimit, lr *limiting.Limiter, key string, blocking bool) {
	if blocking {
		release, err := lr.Acquire(context.Background(), key)
		if err != nil {
			tc.failCount.Add(1)
			return
		}
		defer release()
	} else {
		release, ok := lr.TryAcquire(key)
		if !ok {
			tc.failCount.Add(1)
			return
		}
		defer release()
	}
	time.Sleep(operationDuration)
}

func executeTasks(tasks ...func()) {
	var g sync.WaitGroup

	g.Add(len(tasks))
	for _, task := range tasks {
		go func() {
			defer g.Done()
			task()
		}()
	}
	g.Wait()
}

func executeTaskN(N int, task func()) {
	tasks := make([]func(), N)
	for i := range N {
		tasks[i] = task
	}
	executeTasks(tasks...)
}

func verifyResults(t *testing.T, testCases []*testKeyLimit, blocking bool) {
	for _, tc := range testCases {
		var expectedFailCount int
		if blocking || tc.withoutLimit {
			expectedFailCount = 0
		} else {
			expectedFailCount = max(operationCount*len(tc.keys)-tc.limit, 0)
		}
		actualFailCount := int(tc.failCount.Load())
		require.Equal(t, expectedFailCount, actualFailCount)
	}
}

func resetFailCounts(testLimits []*testKeyLimit) {
	for _, tc := range testLimits {
		tc.failCount.Store(0)
	}
}