package resourcelimiter

import (
	"context"
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/scheduler/nodemanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/testutils"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
)

func TestResourceLimiterOperations(t *testing.T) {
	// Initialize mock pool manager.
	mockPoolConfig := testutils.DefaultConfig()
	var testPool1 rm.WorkerPool = "test-pool-1"
	testPool1Mem := resource.MustParse("1Gi")
	testPool1CPU := resource.MustParse("1")
	var testPool2 rm.WorkerPool = "test-pool-2"
	testPool2Mem := resource.MustParse("2Gi")
	testPool2CPU := resource.MustParse("2")
	mockPoolConfig.PoolsConfig = testutils.PoolsConfig{
		Pools: map[rm.WorkerPool]*testutils.PoolConfig{
			testPool1: {Resources: v1.ResourceRequirements{Limits: v1.ResourceList{
				v1.ResourceCPU:    testPool1CPU,
				v1.ResourceMemory: testPool1Mem,
			}}},
			testPool2: {Resources: v1.ResourceRequirements{Limits: v1.ResourceList{
				v1.ResourceCPU:    testPool2CPU,
				v1.ResourceMemory: testPool2Mem,
			}}},
		},
	}
	mockPoolManager, err := testutils.NewMockManager(mockPoolConfig)
	assert.NoError(t, err)
	pool, err := mockPoolManager.GetPools()
	assert.NoError(t, err)
	assert.NotZero(t, pool)
	limiter := NewResourceLimiter(mockPoolManager)

	// Test update limit.
	err = limiter.UpdateLimit("200Gi", "100", "20Gi", "10")
	assert.NoError(t, err)
	expectedTotalMemory := resource.MustParse("200Gi")
	assert.Equal(t, expectedTotalMemory, limiter.(*resourceLimiter).limit.memTotal)
	expectedTotalCPU := resource.MustParse("100")
	assert.Equal(t, expectedTotalCPU, limiter.(*resourceLimiter).limit.cpuTotal)
	expectedKeyspaceMemory := resource.MustParse("20Gi")
	assert.Equal(t, expectedKeyspaceMemory, limiter.(*resourceLimiter).limit.memKeyspace)
	expectedKeyspaceCPU := resource.MustParse("10")
	assert.Equal(t, expectedKeyspaceCPU, limiter.(*resourceLimiter).limit.cpuKeyspace)

	var largeBatch []string
	for i := 0; i < 100; i++ {
		largeBatch = append(largeBatch, "batch_task_"+strconv.Itoa(i))
	}

	t.Run("test acquire", func(t *testing.T) {
		type tc struct {
			pool             rm.WorkerPool
			keyspace         string
			taskIDs          []string
			success          bool
			perTaskMem       resource.Quantity
			perTaskCPU       resource.Quantity
			totalMemAfter    resource.Quantity
			totalCPUAfter    resource.Quantity
			keyspaceMemAfter resource.Quantity
			keyspaceCPUAfter resource.Quantity
		}
		testCases := []tc{
			{ // Test acquire 1.
				pool:             testPool1,
				keyspace:         "keyspace_1",
				taskIDs:          []string{"task_1"},
				success:          true,
				perTaskMem:       testPool1Mem,
				perTaskCPU:       testPool1CPU,
				totalMemAfter:    *resource.NewQuantity(testPool1Mem.Value()*1, testPool1Mem.Format),
				totalCPUAfter:    *resource.NewQuantity(testPool1CPU.Value()*1, testPool1CPU.Format),
				keyspaceMemAfter: *resource.NewQuantity(testPool1Mem.Value()*1, testPool1Mem.Format),
				keyspaceCPUAfter: *resource.NewQuantity(testPool1CPU.Value()*1, testPool1CPU.Format),
			},
			{ // Test acquire multiple.
				pool:             testPool1,
				keyspace:         "keyspace_1",
				taskIDs:          []string{"task_2", "task_3", "task_4"},
				success:          true,
				perTaskMem:       testPool1Mem,
				perTaskCPU:       testPool1CPU,
				totalMemAfter:    *resource.NewQuantity(testPool1Mem.Value()*4, testPool1Mem.Format),
				totalCPUAfter:    *resource.NewQuantity(testPool1CPU.Value()*4, testPool1CPU.Format),
				keyspaceMemAfter: *resource.NewQuantity(testPool1Mem.Value()*4, testPool1Mem.Format),
				keyspaceCPUAfter: *resource.NewQuantity(testPool1CPU.Value()*4, testPool1CPU.Format),
			},
			{ // Different counter for different keyspace.
				pool:             testPool1,
				keyspace:         "keyspace_2",
				taskIDs:          []string{"task_5"},
				success:          true,
				perTaskMem:       testPool1Mem,
				perTaskCPU:       testPool1CPU,
				totalMemAfter:    *resource.NewQuantity(testPool1Mem.Value()*5, testPool1Mem.Format),
				totalCPUAfter:    *resource.NewQuantity(testPool1CPU.Value()*5, testPool1CPU.Format),
				keyspaceMemAfter: *resource.NewQuantity(testPool1Mem.Value()*1, testPool1Mem.Format),
				keyspaceCPUAfter: *resource.NewQuantity(testPool1CPU.Value()*1, testPool1CPU.Format),
			},
			{ // Error when acquire same task for the same keyspace.
				pool:     testPool1,
				keyspace: "keyspace_2",
				taskIDs:  []string{"task_5"},
				success:  false,
			},
			{ // Error when beyond keyspace limit.
				pool:     testPool1,
				keyspace: "keyspace_2",
				taskIDs:  largeBatch,
				success:  false,
			},
			{ // Failed acquisition does not affect recorded usage.
				pool:             testPool1,
				keyspace:         "keyspace_1",
				success:          true,
				perTaskMem:       testPool1Mem,
				perTaskCPU:       testPool1CPU,
				totalMemAfter:    *resource.NewQuantity(testPool1Mem.Value()*5, testPool1Mem.Format),
				totalCPUAfter:    *resource.NewQuantity(testPool1CPU.Value()*5, testPool1CPU.Format),
				keyspaceMemAfter: *resource.NewQuantity(testPool1Mem.Value()*4, testPool1Mem.Format),
				keyspaceCPUAfter: *resource.NewQuantity(testPool1CPU.Value()*4, testPool1CPU.Format),
			},
			{ // Acquire from different pool is also ok.
				pool:             testPool2,
				keyspace:         "keyspace_1",
				success:          true,
				taskIDs:          []string{"task_6"},
				perTaskMem:       testPool2Mem,
				perTaskCPU:       testPool2CPU,
				totalMemAfter:    *resource.NewQuantity(testPool1Mem.Value()*5+testPool2Mem.Value(), testPool1Mem.Format),
				totalCPUAfter:    *resource.NewQuantity(testPool1CPU.Value()*5+testPool2CPU.Value(), testPool1CPU.Format),
				keyspaceMemAfter: *resource.NewQuantity(testPool1Mem.Value()*4+testPool2Mem.Value(), testPool1Mem.Format),
				keyspaceCPUAfter: *resource.NewQuantity(testPool1CPU.Value()*4+testPool2CPU.Value(), testPool1CPU.Format),
			},
		}
		for _, testCase := range testCases {
			err = limiter.Acquire(testCase.keyspace, testCase.taskIDs, testCase.pool)
			if !testCase.success {
				assert.Error(t, err)
				continue
			}
			assert.NoError(t, err)
			// verify total usage.
			assert.True(t, testCase.totalMemAfter.Equal(limiter.(*resourceLimiter).mem))
			assert.True(t, testCase.totalCPUAfter.Equal(limiter.(*resourceLimiter).cpu))
			// verify keyspace usage.
			assert.True(t, testCase.keyspaceMemAfter.Equal(limiter.(*resourceLimiter).usageRecords[testCase.keyspace].mem))
			assert.True(t, testCase.keyspaceCPUAfter.Equal(limiter.(*resourceLimiter).usageRecords[testCase.keyspace].cpu))
			// verify each task.
			for _, taskID := range testCase.taskIDs {
				assert.True(t, testCase.perTaskMem.Equal(limiter.(*resourceLimiter).usageRecords[testCase.keyspace].
					tasks[taskID].mem))
				assert.True(t, testCase.perTaskCPU.Equal(limiter.(*resourceLimiter).usageRecords[testCase.keyspace].
					tasks[taskID].cpu))
			}
		}
	})

	t.Run("test release", func(t *testing.T) {
		type tc struct {
			keyspace         string
			taskID           string
			success          bool
			totalMemAfter    resource.Quantity
			totalCPUAfter    resource.Quantity
			keyspaceMemAfter resource.Quantity
			keyspaceCPUAfter resource.Quantity
		}
		testCases := []tc{
			{
				// Acquire from different pool is also ok.
				keyspace:         "keyspace_1",
				success:          true,
				taskID:           "task_6",
				totalMemAfter:    *resource.NewQuantity(testPool1Mem.Value()*5, testPool1Mem.Format),
				totalCPUAfter:    *resource.NewQuantity(testPool1CPU.Value()*5, testPool1CPU.Format),
				keyspaceMemAfter: *resource.NewQuantity(testPool1Mem.Value()*4, testPool1Mem.Format),
				keyspaceCPUAfter: *resource.NewQuantity(testPool1CPU.Value()*4, testPool1CPU.Format),
			},
			{ // Release a non-existing task would fail.
				keyspace: "keyspace_1",
				success:  false,
				taskID:   "task_6",
			},
			{ // Release a non-existing keyspace would fail.
				keyspace: "keyspace_xxx",
				success:  false,
			},
		}

		for _, testCase := range testCases {
			err = limiter.Release(testCase.keyspace, testCase.taskID)
			if !testCase.success {
				assert.Error(t, err)
				continue
			}
			assert.NoError(t, err)
			// verify total usage.
			assert.True(t, testCase.totalMemAfter.Equal(limiter.(*resourceLimiter).mem))
			assert.True(t, testCase.totalCPUAfter.Equal(limiter.(*resourceLimiter).cpu))
			// verify keyspace usage.
			assert.True(t, testCase.keyspaceMemAfter.Equal(limiter.(*resourceLimiter).usageRecords[testCase.keyspace].mem))
			assert.True(t, testCase.keyspaceCPUAfter.Equal(limiter.(*resourceLimiter).usageRecords[testCase.keyspace].cpu))
			// verify task gone.
			assert.NotContains(t, limiter.(*resourceLimiter).usageRecords[testCase.keyspace].tasks, testCase.taskID)
		}
	})

	// update to high limit and acquire large batch.
	err = limiter.UpdateLimit("999999Gi", "999999", "999999Gi", "999999")
	assert.NoError(t, err)
	err = limiter.Acquire("keyspace_1", largeBatch, testPool1)
	assert.NoError(t, err)
}

func TestResourceManagerInitialization(t *testing.T) {
	// Initialize mock pool manager.
	mockPoolConfig := testutils.DefaultConfig()
	var testPool rm.WorkerPool = "test-pool-1"
	testPoolMem := resource.MustParse("1Gi")
	testPoolCPU := resource.MustParse("1")
	mockPoolConfig.PoolsConfig = testutils.PoolsConfig{
		Pools: map[rm.WorkerPool]*testutils.PoolConfig{
			testPool: {Resources: v1.ResourceRequirements{Limits: v1.ResourceList{
				v1.ResourceCPU:    testPoolCPU,
				v1.ResourceMemory: testPoolMem,
			}}},
		},
	}
	mockPoolManager, err := testutils.NewMockManager(mockPoolConfig)
	assert.NoError(t, err)
	pool, err := mockPoolManager.GetPools()
	assert.NoError(t, err)
	assert.NotZero(t, pool)
	limiter := NewResourceLimiter(mockPoolManager)
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	assert.NoError(t, err)
	err = db.AutoMigrate(&nodemanager.Node{})
	assert.NoError(t, err)

	err = db.WithContext(context.Background()).Create(&nodemanager.Node{
		NodeID:          "test-node-1",
		TaskID:          "test-task-1",
		KeyspaceName:    "test-keyspace-1",
		PoolName:        string(testPool),
		ActivateRequest: "{}",
		IsBound:         true,
	}).Error
	assert.NoError(t, err)
	err = db.WithContext(context.Background()).Create(&nodemanager.Node{
		NodeID:          "test-node-2",
		TaskID:          "test-task-2",
		KeyspaceName:    "test-keyspace-1",
		PoolName:        string(testPool),
		ActivateRequest: "{}",
		IsBound:         true,
	}).Error
	assert.NoError(t, err)
	err = db.WithContext(context.Background()).Create(&nodemanager.Node{
		NodeID:          "test-node-3",
		TaskID:          "test-task-3",
		KeyspaceName:    "test-keyspace-2",
		PoolName:        string(testPool),
		ActivateRequest: "{}",
		IsBound:         true,
	}).Error
	assert.NoError(t, err)

	mockNodeManager := nodemanager.NewMockNodeManager(mockPoolManager, db)
	err = limiter.Initialize(context.Background(), mockNodeManager)
	assert.NoError(t, err)

	// Total usage should be 3 tasks.
	expectedTotalMem := *resource.NewQuantity(testPoolMem.Value()*3, testPoolMem.Format)
	expectedTotalCPU := *resource.NewQuantity(testPoolCPU.Value()*3, testPoolCPU.Format)
	assert.True(t, limiter.(*resourceLimiter).mem.Equal(expectedTotalMem))
	assert.True(t, limiter.(*resourceLimiter).cpu.Equal(expectedTotalCPU))

	// Keyspace 1 usage should be 2 tasks.
	expectedKeyspace1Mem := *resource.NewQuantity(testPoolMem.Value()*2, testPoolMem.Format)
	expectedKeyspace1CPU := *resource.NewQuantity(testPoolCPU.Value()*2, testPoolCPU.Format)
	assert.True(t, limiter.(*resourceLimiter).usageRecords["test-keyspace-1"].mem.Equal(expectedKeyspace1Mem))
	assert.True(t, limiter.(*resourceLimiter).usageRecords["test-keyspace-1"].cpu.Equal(expectedKeyspace1CPU))

	// Keyspace 2 usage should be 1 task.
	expectedKeyspace2Mem := *resource.NewQuantity(testPoolMem.Value()*1, testPoolMem.Format)
	expectedKeyspace2CPU := *resource.NewQuantity(testPoolCPU.Value()*1, testPoolCPU.Format)
	assert.True(t, limiter.(*resourceLimiter).usageRecords["test-keyspace-2"].mem.Equal(expectedKeyspace2Mem))
	assert.True(t, limiter.(*resourceLimiter).usageRecords["test-keyspace-2"].cpu.Equal(expectedKeyspace2CPU))
}