package nodemanager

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/job"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/poolmanager/factory"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/testutils"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

const (
	sleepTime   = 15 * time.Second
	defaultPool = "worker-100m-128Mi"
)

func startMockNodeManager(t *testing.T, nm *manager) (context.CancelFunc, <-chan error) {
	t.Helper()

	ctx, cancel := context.WithCancel(context.Background())
	errCh := make(chan error, 1)
	go func() {
		errCh <- nm.Start(ctx)
	}()
	return cancel, errCh
}

func TestNodeManagerAssignAndRecycleNode(t *testing.T) {
	ctx := context.Background()

	dbPath := filepath.Join(t.TempDir(), "TestNodeManagerAssignAndRecycleNode.db")
	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
	assert.NoError(t, err)
	defer func() {
		err := os.Remove(dbPath)
		assert.NoError(t, err)
	}()

	testKeyspaceName := "ks1"
	taskID := "taskID_01"
	testPoolName := defaultPool

	mockConfig := testutils.DefaultConfig()
	mockManager, err := testutils.NewMockManager(mockConfig)
	//	Get all pool size
	var allPoolSize int
	var testPoolSize int
	for poolName, poolConfig := range mockConfig.PoolsConfig.Pools {
		if testPoolName == string(poolName) {
			testPoolSize += poolConfig.Size
		}
		allPoolSize += poolConfig.Size
	}

	nm := NewMockNodeManager(mockManager, db)
	_ = nm.dbAutoMigrate(ctx)
	err = nm.UpdateNodeList(ctx)
	assert.NoError(t, err)

	time.Sleep(sleepTime)
	// Check all nodes was initialized.
	locks, err := nm.nodeBinder.GetAllNodes(ctx)
	assert.NoError(t, err)
	assert.Equal(t, allPoolSize, len(locks))

	// Assign worker node
	activateReq := tidb.ActivateRequest{
		KeyspaceName: testKeyspaceName,
		WorkerRole:   string(job.TypeAutoAnalyze),
	}
	tasks := []*Request{{
		ActivateRequest: activateReq,
		PoolName:        defaultPool,
		TaskID:          taskID,
	},
	}
	err = nm.Assign(ctx, testKeyspaceName, tasks)
	assert.NoError(t, err)

	// Get assign msg
	lockRes, err := nm.nodeBinder.GetNodeByTaskID(ctx, taskID)
	assert.NoError(t, err)
	assert.True(t, strings.HasPrefix(lockRes.NodeID, factory.PoolManagerPrefix+"-"+testPoolName))
	assert.Equal(t, taskID, lockRes.TaskID)

	checkNode, err := nm.nodeBinder.GetNodeByNodeID(ctx, lockRes.NodeID)
	assert.NoError(t, err)
	assert.Equal(t, checkNode.TaskID, taskID)

	// Get standby workers, the standby sets should not contain IsBound node.
	workers, err := nm.PoolManager.GetStandbyWorkers(rm.WorkerPool(testPoolName))
	assert.NoError(t, err)
	assert.Equal(t, testPoolSize-1, len(workers))
	for _, worker := range workers {
		assert.NotEqual(t, lockRes.TaskID, worker)
	}

	// Recycle worker node
	err = nm.Recycle(ctx, testKeyspaceName, taskID)
	assert.NoError(t, err)

	// Check that ksTaskID does not have a corresponding lock
	lockRes, err = nm.nodeBinder.GetNodeByTaskID(ctx, taskID)
	assert.True(t, strings.Contains(err.Error(), "record not found"))
	assert.Nil(t, nil, lockRes)

	// Get standby workers, all workers ars standby.
	workers, err = nm.PoolManager.GetStandbyWorkers(rm.WorkerPool(testPoolName))
	assert.NoError(t, err)
	assert.Equal(t, len(workers), testPoolSize)
}

// Check pool manager and locks data
func TestNodeManagerLoop(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "TestNodeManagerLoop.db")
	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
	assert.NoError(t, err)
	defer func() {
		err := os.Remove(dbPath)
		assert.NoError(t, err)
	}()

	testPoolName := defaultPool
	mockConfig := testutils.DefaultConfig()
	mockManager, err := testutils.NewMockManager(mockConfig)
	assert.NoError(t, err)

	//	Get all pool size
	var allPoolSize int
	var testPoolSize int
	for poolName, poolConfig := range mockConfig.PoolsConfig.Pools {
		if testPoolName == string(poolName) {
			testPoolSize += poolConfig.Size
		}
		allPoolSize += poolConfig.Size
	}

	nm := NewMockNodeManager(mockManager, db)

	cancel, errCh := startMockNodeManager(t, nm)
	defer func() {
		cancel()
		assert.ErrorIs(t, <-errCh, context.Canceled)
	}()

	ctx := context.Background()

	// Check updateNodeListLoop
	time.Sleep(sleepTime)
	locks, err := nm.nodeBinder.GetAllNodes(ctx)
	assert.NoError(t, err)
	assert.Equal(t, allPoolSize, len(locks))

	// Check reactivateRestartNodeLoop ---------------------------------------------
	testKeyspaceName := "ks1"
	taskID := "taskID_01"

	// active task
	// Assign worker node
	activateReq := tidb.ActivateRequest{
		KeyspaceName: testKeyspaceName,
		WorkerRole:   string(job.TypeAutoAnalyze),
	}
	tasks := []*Request{{
		ActivateRequest: activateReq,
		PoolName:        defaultPool,
		TaskID:          taskID,
	},
	}
	err = nm.Assign(ctx, testKeyspaceName, tasks)
	assert.NoError(t, err)

	time.Sleep(sleepTime)
	workers, err := nm.PoolManager.GetStandbyWorkers(rm.WorkerPool(testPoolName))
	assert.NoError(t, err)
	assert.Equal(t, testPoolSize-1, len(workers))

	mockConfig2 := testutils.DefaultConfig()
	mockManager2, err := testutils.NewMockManager(mockConfig2)
	assert.NoError(t, err)

	cancel()
	assert.ErrorIs(t, <-errCh, context.Canceled)

	nm = NewMockNodeManager(mockManager2, db)
	cancel, errCh = startMockNodeManager(t, nm)

	workers, err = nm.PoolManager.GetStandbyWorkers(rm.WorkerPool(testPoolName))
	assert.NoError(t, err)
	assert.Equal(t, testPoolSize, len(workers))

	// wait reactivateRestartNodeLoop activate standby pool
	time.Sleep(sleepTime)

	workers, err = nm.PoolManager.GetStandbyWorkers(rm.WorkerPool(testPoolName))
	assert.NoError(t, err)
	assert.Equal(t, testPoolSize-1, len(workers))
}