package encryptionkey

import (
	"context"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/encryptionchange"
	mockexecutor "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/encryptionchange/mock_encryptionchange"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/workflow/action"
	"github.com/tidbcloud/tidb-management-service/utilities/idgenerator"
)

//nolint:contextcheck // Run() takes run.ExecuteContext which carries ctx internally
func TestUpdateClusterEncryptionActions(t *testing.T) {
	ctx := t.Context()
	trans := *dbTransactioner.Load()
	sfid := idgenerator.NewSonyflakeID(2, "")

	ts := newTestService(t)
	defer ts.ctrl.Finish()

	t.Run("GenMeta/Success", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().GenMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.GenMetaResponse{
				FileId:    100,
				PdUpdated: true,
				PdAfter:   encryptionchange.PdAfterState{CurrentFileId: 100, NextFileId: 0},
				EncryptionConfig: encryptionchange.EncryptionConfigInfo{
					Enabled: true,
					CmekId:  "test-cmek",
				},
			}, nil)

		act := action.NewRunCseCtlGenMetaAction(ts.clusterRepo, ts.regionRepo, mockExec, ts.encryptionKeyRepo)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.NoError(t, err)
	})

	t.Run("GenMeta/ExecutorFailure", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().GenMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil, fmt.Errorf("kms unreachable"))

		act := action.NewRunCseCtlGenMetaAction(ts.clusterRepo, ts.regionRepo, mockExec, ts.encryptionKeyRepo)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "kms unreachable")
	})

	t.Run("SyncMeta/Success", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().SyncMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.SyncMetaResponse{Started: true, Running: true}, nil)

		act := action.NewRunCseCtlSyncMetaAction(ts.clusterRepo, ts.regionRepo, mockExec)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.NoError(t, err)
	})

	t.Run("SyncMeta/NotStarted", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().SyncMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.SyncMetaResponse{Started: false, Running: false}, nil)

		act := action.NewRunCseCtlSyncMetaAction(ts.clusterRepo, ts.regionRepo, mockExec)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "sync-meta returned not-started/not-running")
	})

	t.Run("PollStatus/Converged", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().GetStatus(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.StatusResponse{
				EncryptionEnabled: true,
				MetaState:         encryptionchange.MetaStateInfo{CurrentFileId: 100, NextFileId: 0},
			}, nil)

		act := action.NewPollCseCtlStatusAction(ts.clusterRepo, ts.regionRepo, mockExec)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.NoError(t, err)
	})

	t.Run("PollStatus/NotConverged", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().GetStatus(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.StatusResponse{
				EncryptionEnabled: false,
				MetaState:         encryptionchange.MetaStateInfo{CurrentFileId: 100, NextFileId: 50},
			}, nil).
			Times(2)
		mockExec.EXPECT().SyncMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(_ context.Context, _ *entity.Region, req *encryptionchange.SyncMetaRequest) (*encryptionchange.SyncMetaResponse, error) {
				assert.Equal(t, fmt.Sprint(cluster.GetClusterId()), req.ClusterId)
				assert.Equal(t, cluster.GetSharedPoolId(), req.SharedPoolId)
				assert.Equal(t, fmt.Sprint(cluster.GetClusterId()), req.KeyspaceName)
				return &encryptionchange.SyncMetaResponse{Started: true, Running: true}, nil
			}).
			Times(2)

		act := action.NewPollCseCtlStatusAction(ts.clusterRepo, ts.regionRepo, mockExec)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		// Should return a retryable error (not converged yet)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "not converged")

		err = act.Run(execCtx, nil)
		// sync-meta is safe and expected to run again on a repeated non-converged poll.
		require.Error(t, err)
		assert.Contains(t, err.Error(), "not converged")
	})

	t.Run("PollStatus/NotConvergedSyncMetaFailure", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().GetStatus(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.StatusResponse{
				EncryptionEnabled: false,
				MetaState:         encryptionchange.MetaStateInfo{CurrentFileId: 100, NextFileId: 50},
			}, nil)
		mockExec.EXPECT().SyncMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil, fmt.Errorf("sync-meta failed"))

		act := action.NewPollCseCtlStatusAction(ts.clusterRepo, ts.regionRepo, mockExec)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "sync-meta failed")
	})

	t.Run("PollStatus/NotConvergedSyncMetaNotStarted", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)

		mockExec.EXPECT().GetStatus(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.StatusResponse{
				EncryptionEnabled: false,
				MetaState:         encryptionchange.MetaStateInfo{CurrentFileId: 100, NextFileId: 50},
			}, nil)
		mockExec.EXPECT().SyncMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.SyncMetaResponse{Started: false, Running: false}, nil)

		act := action.NewPollCseCtlStatusAction(ts.clusterRepo, ts.regionRepo, mockExec)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "sync-meta retry returned not-started/not-running")
	})

	t.Run("Finalize/CMEK", func(t *testing.T) {
		cleanTestData(ctx, t)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		newMasterKey := "arn:aws:kms:us-west-2:123456789012:key/new-cmek-key"
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.MasterKey = newMasterKey
		})
		createTestClusterEncryptionBinding(ctx, t, ts.bindingRepo,
			cluster.GetClusterId(), key.Id,
			entity.ClusterEncryptionBindingStateModifying,
		)

		act := action.NewFinalizeEncryptionChangeAction(ts.clusterRepo, ts.bindingRepo, ts.encryptionKeyRepo, trans)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.NoError(t, err)

		// Key should be Active with updated master key
		loadedKey, keyErr := ts.encryptionKeyRepo.Load(ctx, key.Id, false)
		require.Nil(t, keyErr)
		assert.Equal(t, newMasterKey, loadedKey.MasterKey)

		// Cluster should reference the new key
		loadedCluster, clusterErr := ts.clusterRepo.Get(ctx, cluster.GetClusterId(), false)
		require.Nil(t, clusterErr)
		assert.Equal(t, key.Id, loadedCluster.EncryptionKeyId)

		// Binding should be Enabled
		loadedBinding, bindErr := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		require.Nil(t, bindErr)
		assert.Equal(t, entity.ClusterEncryptionBindingStateEnabled, loadedBinding.State)
	})

	t.Run("Finalize/DefaultKey", func(t *testing.T) {
		cleanTestData(ctx, t)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		originalMasterKey := "arn:aws:kms:us-west-2:123456789012:key/original-key"
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.MasterKey = originalMasterKey
			k.CustomerManaged = false
		})
		createTestClusterEncryptionBinding(ctx, t, ts.bindingRepo,
			cluster.GetClusterId(), key.Id,
			entity.ClusterEncryptionBindingStateModifying,
		)

		act := action.NewFinalizeEncryptionChangeAction(ts.clusterRepo, ts.bindingRepo, ts.encryptionKeyRepo, trans)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-default-key", "enable-default-key")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.NoError(t, err)

		// Key should be Active with unchanged master key (default_key doesn't update it)
		loadedKey, keyErr := ts.encryptionKeyRepo.Load(ctx, key.Id, false)
		require.Nil(t, keyErr)
		assert.Equal(t, originalMasterKey, loadedKey.MasterKey)

		// Cluster should reference the key
		loadedCluster, clusterErr := ts.clusterRepo.Get(ctx, cluster.GetClusterId(), false)
		require.Nil(t, clusterErr)
		assert.Equal(t, key.Id, loadedCluster.EncryptionKeyId)

		// Binding should be Enabled
		loadedBinding, bindErr := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		require.Nil(t, bindErr)
		assert.Equal(t, entity.ClusterEncryptionBindingStateEnabled, loadedBinding.State)
	})

	t.Run("Finalize/DisableCMEK", func(t *testing.T) {
		cleanTestData(ctx, t)

		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.EncryptionKeyId = key.Id
		})
		createTestClusterEncryptionBinding(ctx, t, ts.bindingRepo,
			cluster.GetClusterId(), key.Id,
			entity.ClusterEncryptionBindingStateModifying,
		)

		act := action.NewFinalizeEncryptionChangeAction(ts.clusterRepo, ts.bindingRepo, ts.encryptionKeyRepo, trans)
		vars := workflowVars(cluster.GetClusterId(), key.Id, "disable-cmek", "disable-cmek")
		execCtx := newTestExecuteContext(ctx, vars)

		err := act.Run(execCtx, nil)
		require.NoError(t, err)

		// Cluster should have EncryptionKeyId = 0
		loadedCluster, clusterErr := ts.clusterRepo.Get(ctx, cluster.GetClusterId(), false)
		require.Nil(t, clusterErr)
		assert.Equal(t, uint64(0), loadedCluster.EncryptionKeyId)

		// Binding should be deleted
		_, bindErr := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		assert.NotNil(t, bindErr)
	})

	t.Run("EndToEnd_AllActions", func(t *testing.T) {
		cleanTestData(ctx, t)
		ctrl := gomock.NewController(t)
		mockExec := mockexecutor.NewMockExecutor(ctrl)

		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		targetKey := "arn:aws:kms:us-west-2:123456789012:key/e2e-cmek-key"
		key := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.MasterKey = targetKey
		})
		createTestClusterEncryptionBinding(ctx, t, ts.bindingRepo,
			cluster.GetClusterId(), key.Id,
			entity.ClusterEncryptionBindingStateModifying,
		)

		vars := workflowVars(cluster.GetClusterId(), key.Id, "enable-cmek", "enable-cmek")

		// Mock all executor calls in sequence
		mockExec.EXPECT().GenMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.GenMetaResponse{
				FileId:    1,
				PdUpdated: true,
				PdAfter:   encryptionchange.PdAfterState{CurrentFileId: 1, NextFileId: 0},
			}, nil)
		mockExec.EXPECT().SyncMeta(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.SyncMetaResponse{Started: true, Running: true}, nil)
		mockExec.EXPECT().GetStatus(gomock.Any(), gomock.Any(), gomock.Any()).
			Return(&encryptionchange.StatusResponse{
				EncryptionEnabled: true,
				MetaState:         encryptionchange.MetaStateInfo{CurrentFileId: 1, NextFileId: 0},
			}, nil)

		// 1. GenMeta
		genMetaAct := action.NewRunCseCtlGenMetaAction(ts.clusterRepo, ts.regionRepo, mockExec, ts.encryptionKeyRepo)
		err := genMetaAct.Run(newTestExecuteContext(ctx, vars), nil)
		require.NoError(t, err)

		// 2. SyncMeta
		syncMetaAct := action.NewRunCseCtlSyncMetaAction(ts.clusterRepo, ts.regionRepo, mockExec)
		err = syncMetaAct.Run(newTestExecuteContext(ctx, vars), nil)
		require.NoError(t, err)

		// 3. PollStatus
		pollAct := action.NewPollCseCtlStatusAction(ts.clusterRepo, ts.regionRepo, mockExec)
		err = pollAct.Run(newTestExecuteContext(ctx, vars), nil)
		require.NoError(t, err)

		// 4. Finalize
		finalizeAct := action.NewFinalizeEncryptionChangeAction(ts.clusterRepo, ts.bindingRepo, ts.encryptionKeyRepo, trans)
		err = finalizeAct.Run(newTestExecuteContext(ctx, vars), nil)
		require.NoError(t, err)

		// Verify final state
		loadedKey, keyErr := ts.encryptionKeyRepo.Load(ctx, key.Id, false)
		require.Nil(t, keyErr)
		assert.Equal(t, targetKey, loadedKey.MasterKey)

		loadedCluster, clusterErr := ts.clusterRepo.Get(ctx, cluster.GetClusterId(), false)
		require.Nil(t, clusterErr)
		assert.Equal(t, key.Id, loadedCluster.EncryptionKeyId)

		loadedBinding, bindErr := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		require.Nil(t, bindErr)
		assert.Equal(t, entity.ClusterEncryptionBindingStateEnabled, loadedBinding.State)
	})
}