package standbywatcher

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"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"
)

func mustNewStandbyWatcher(re *require.Assertions) (rm.StandbyWatcher, testutils.MockManager) {
	mockConfig := testutils.DefaultConfig()
	mockManager, err := testutils.NewMockManager(mockConfig)
	re.NoError(err)

	return NewStandbyWatcher(mockManager, &rm.StandbyWatcherConfig{
		ProbeInterval: 300 * time.Millisecond,
	}), mockManager
}

func TestStandbyWatcher(t *testing.T) {
	re := require.New(t)
	watcher, manager := mustNewStandbyWatcher(re)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		re.NoError(watcher.Start(ctx))
	}()

	cfg := manager.GetConfig()
	// All workers are in initial standby state.
	re.EventuallyWithT(func(c *assert.CollectT) {
		for pool := range cfg.PoolsConfig.Pools {
			expectedStandbyWorkers, err := manager.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			actualStandbyWorkers, err := watcher.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			assert.ElementsMatch(c, expectedStandbyWorkers, actualStandbyWorkers)
			assert.Len(c, actualStandbyWorkers, cfg.PoolsConfig.Pools[pool].Size)
		}
	}, 3*time.Second, 500*time.Millisecond)

	// Remove a pool from the configmap, and check if the watcher can detect change.
	var deletedPool rm.WorkerPool
	for pool := range cfg.PoolsConfig.Pools {
		delete(cfg.PoolsConfig.Pools, pool)
		deletedPool = pool
		break
	}
	re.NoError(manager.UpdateCluster(ctx, cfg))
	re.EventuallyWithT(func(c *assert.CollectT) {
		for pool := range cfg.PoolsConfig.Pools {
			expectedStandbyWorkers, err := manager.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			actualStandbyWorkers, err := watcher.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			assert.ElementsMatch(c, expectedStandbyWorkers, actualStandbyWorkers)
			assert.Len(c, actualStandbyWorkers, cfg.PoolsConfig.Pools[pool].Size)
		}
		deletedPoolStandbyWorkers, err := manager.GetStandbyWorkers(deletedPool)
		assert.NoError(c, err)
		assert.Empty(c, deletedPoolStandbyWorkers)
	}, 3*time.Second, 500*time.Millisecond)

	// Activate one worker in the pool, and check if the watcher can detect the change.
	var (
		activatedPool   rm.WorkerPool
		activatedWorker rm.Worker
	)
	for pool := range cfg.PoolsConfig.Pools {
		standbyWorkers, err := manager.GetStandbyWorkers(pool)
		assert.NoError(t, err)
		if len(standbyWorkers) > 0 {
			activatedWorker = standbyWorkers[0]
			activatedPool = pool
			break
		}
	}
	re.NotEmpty(activatedWorker)
	re.NoError(manager.ActivateWorker(
		context.Background(),
		activatedPool,
		activatedWorker,
		tidb.ActivateRequest{},
	))
	re.EventuallyWithT(func(c *assert.CollectT) {
		for pool := range cfg.PoolsConfig.Pools {
			expectedStandbyWorkers, err := manager.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			actualStandbyWorkers, err := watcher.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			assert.ElementsMatch(c, expectedStandbyWorkers, actualStandbyWorkers)
			expectedSize := cfg.PoolsConfig.Pools[pool].Size
			if pool == activatedPool {
				expectedSize--
			}
			assert.Len(c, actualStandbyWorkers, expectedSize)
		}
	}, 3*time.Second, 500*time.Millisecond)
	activatedPoolStandbyWorkers, err := manager.GetStandbyWorkers(activatedPool)
	re.NoError(err)
	re.NotContains(activatedPoolStandbyWorkers, activatedWorker)

	// Recycle the activated worker, and check if the watcher can detect the change.
	re.NoError(manager.RecycleWorker(context.Background(), activatedPool, activatedWorker, ""))
	re.EventuallyWithT(func(c *assert.CollectT) {
		for pool := range cfg.PoolsConfig.Pools {
			expectedStandbyWorkers, err := manager.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			actualStandbyWorkers, err := watcher.GetStandbyWorkers(pool)
			assert.NoError(c, err)
			assert.ElementsMatch(c, expectedStandbyWorkers, actualStandbyWorkers)
			assert.Len(c, actualStandbyWorkers, cfg.PoolsConfig.Pools[pool].Size)
		}
	}, 3*time.Second, 500*time.Millisecond)
	activatedPoolStandbyWorkers, err = manager.GetStandbyWorkers(activatedPool)
	re.NoError(err)
	re.Contains(activatedPoolStandbyWorkers, activatedWorker)
}