package encryptionkey

import (
	"context"
	"testing"

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

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	regionalserveradapter "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/idgenerator"
)

func TestSwitchStateTransitions(t *testing.T) {
	sfid := idgenerator.NewSonyflakeID(1, "")

	// § 13.1 / §3: state transition validation
	t.Run("AllowedTransitions", func(t *testing.T) {
		tests := []struct {
			name           string
			currentKey     *entity.EncryptionKey
			req            *SwitchEncryptionRequest
			wantTransition string
		}{
			{
				name:       "NoEncryption_to_CMEK",
				currentKey: nil,
				req: &SwitchEncryptionRequest{
					TargetCmek: &globalngv1.CustomerManagedEncryptionKey{
						CloudProviderId: globalcommonv1.CloudProvider_AWS,
					},
				},
				wantTransition: "enable-cmek",
			},
			{
				name:       "NoEncryption_to_DefaultKey",
				currentKey: nil,
				req: &SwitchEncryptionRequest{
					TargetDefaultKey: true,
				},
				wantTransition: "enable-default-key",
			},
			{
				name: "CMEK_to_DefaultKey",
				currentKey: &entity.EncryptionKey{
					Id:              sfid.MustNew(),
					CustomerManaged: true,
					MasterKey:       "arn:aws:kms:us-west-2:123456789012:key/old-key",
				},
				req: &SwitchEncryptionRequest{
					TargetDefaultKey: true,
				},
				wantTransition: "switch-encryption-key",
			},
			{
				name: "DefaultKey_to_CMEK",
				currentKey: &entity.EncryptionKey{
					Id:              sfid.MustNew(),
					CustomerManaged: false,
					MasterKey:       "default_key",
				},
				req: &SwitchEncryptionRequest{
					TargetCmek: &globalngv1.CustomerManagedEncryptionKey{
						CloudProviderId: globalcommonv1.CloudProvider_AWS,
					},
				},
				wantTransition: "switch-encryption-key",
			},
			{
				name: "CMEK_A_to_CMEK_B",
				currentKey: &entity.EncryptionKey{
					Id:              sfid.MustNew(),
					CustomerManaged: true,
					MasterKey:       "arn:aws:kms:us-west-2:123456789012:key/key-a",
				},
				req: &SwitchEncryptionRequest{
					TargetCmek: &globalngv1.CustomerManagedEncryptionKey{
						CloudProviderId: globalcommonv1.CloudProvider_AWS,
					},
				},
				wantTransition: "switch-encryption-key",
			},
			{
				name: "Encrypted_to_Disable",
				currentKey: &entity.EncryptionKey{
					Id:              sfid.MustNew(),
					CustomerManaged: true,
				},
				req: &SwitchEncryptionRequest{
					TargetDisable: true,
				},
				wantTransition: "disable-cmek",
			},
		}

		for _, tc := range tests {
			t.Run(tc.name, func(t *testing.T) {
				got, err := determineTransitionType(tc.currentKey, tc.req)
				require.Nil(t, err)
				assert.Equal(t, tc.wantTransition, got)
			})
		}
	})

	t.Run("DisallowedTransitions", func(t *testing.T) {
		tests := []struct {
			name       string
			currentKey *entity.EncryptionKey
			req        *SwitchEncryptionRequest
		}{
			{
				name:       "NoEncryption_to_NoEncryption",
				currentKey: nil,
				req:        &SwitchEncryptionRequest{},
			},
			{
				name: "Encrypted_to_NoTarget",
				currentKey: &entity.EncryptionKey{
					Id:              sfid.MustNew(),
					CustomerManaged: true,
				},
				req: &SwitchEncryptionRequest{},
			},
		}

		for _, tc := range tests {
			t.Run(tc.name, func(t *testing.T) {
				_, err := determineTransitionType(tc.currentKey, tc.req)
				assert.NotNil(t, err)
			})
		}
	})

	t.Run("NoOpTransitions", func(t *testing.T) {
		// Disabling encryption on an already-unencrypted cluster is a no-op
		ttype, err := determineTransitionType(nil, &SwitchEncryptionRequest{TargetDisable: true})
		assert.Nil(t, err)
		assert.Equal(t, "no-op", ttype)

		// Switching default key to default key is a no-op
		ttype, err = determineTransitionType(&entity.EncryptionKey{
			Id:              sfid.MustNew(),
			CustomerManaged: false,
			MasterKey:       "default_key",
		}, &SwitchEncryptionRequest{TargetDefaultKey: true})
		assert.Nil(t, err)
		assert.Equal(t, "no-op", ttype)

		// Switching CMEK to same CMEK ARN is a no-op
		sameArn := "arn:aws:kms:us-west-2:123456789012:key/same-key"
		ttype, err = determineTransitionType(&entity.EncryptionKey{
			Id:              sfid.MustNew(),
			CustomerManaged: true,
			MasterKey:       sameArn,
		}, &SwitchEncryptionRequest{
			TargetCmek: &globalngv1.CustomerManagedEncryptionKey{
				CloudProviderId: globalcommonv1.CloudProvider_AWS,
				Key: &globalngv1.CustomerManagedEncryptionKey_AwsKms_{
					AwsKms: &globalngv1.CustomerManagedEncryptionKey_AwsKms{
						KmsKeyArn: sameArn,
					},
				},
			},
		})
		assert.Nil(t, err)
		assert.Equal(t, "no-op", ttype)
	})

	t.Run("CseCtlMode", func(t *testing.T) {
		// From no encryption → enable-cmek
		mode := determineCseCtlMode(nil, &SwitchEncryptionRequest{})
		assert.Equal(t, "enable-cmek", mode)

		mode = determineCseCtlMode(&entity.EncryptionKey{Id: 0}, &SwitchEncryptionRequest{})
		assert.Equal(t, "enable-cmek", mode)

		// From existing encryption → also enable-cmek (cse-ctl uses same mode)
		mode = determineCseCtlMode(&entity.EncryptionKey{Id: 123}, &SwitchEncryptionRequest{})
		assert.Equal(t, "enable-cmek", mode)
	})

	t.Run("ResolveTargetMasterKey", func(t *testing.T) {
		// AWS KMS
		awsReq := &SwitchEncryptionRequest{
			TargetCmek: &globalngv1.CustomerManagedEncryptionKey{
				Key: &globalngv1.CustomerManagedEncryptionKey_AwsKms_{
					AwsKms: &globalngv1.CustomerManagedEncryptionKey_AwsKms{
						KmsKeyArn: "arn:aws:kms:us-west-2:123456789012:key/my-key",
					},
				},
			},
		}
		assert.Equal(t, "arn:aws:kms:us-west-2:123456789012:key/my-key", resolveTargetMasterKey(awsReq))

		// Aliyun KMS
		aliyunReq := &SwitchEncryptionRequest{
			TargetCmek: &globalngv1.CustomerManagedEncryptionKey{
				Key: &globalngv1.CustomerManagedEncryptionKey_AliyunKms_{
					AliyunKms: &globalngv1.CustomerManagedEncryptionKey_AliyunKms{
						KmsKeyArn: "acs:kms:ap-southeast-1:12345:key/key-123",
					},
				},
			},
		}
		assert.Equal(t, "acs:kms:ap-southeast-1:12345:key/key-123", resolveTargetMasterKey(aliyunReq))

		// Default key
		defaultReq := &SwitchEncryptionRequest{TargetDefaultKey: true}
		assert.Equal(t, "default_key", resolveTargetMasterKey(defaultReq))
	})
}

func TestSwitchClusterEncryptionKey(t *testing.T) {
	ctx := t.Context()
	sfid := idgenerator.NewSonyflakeID(1, "")

	t.Run("ClusterNotAvailable", func(t *testing.T) {
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.State = globalngv1.Cluster_CREATING
		})

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:    cluster.GetClusterId(),
			TargetCmek:   newTestAwsCmek(),
			ValidateOnly: false,
		})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "AVAILABLE")
	})

	t.Run("ActiveBindingConflict", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)

		// Create a Modifying binding
		createTestClusterEncryptionBinding(ctx, t, ts.bindingRepo,
			cluster.GetClusterId(), sfid.MustNew(),
			entity.ClusterEncryptionBindingStateModifying,
		)

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:    cluster.GetClusterId(),
			TargetCmek:   newTestAwsCmek(),
			ValidateOnly: false,
		})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "already in progress")
	})

	t.Run("ClusterWithoutSharedPoolFailsCMEKVerify", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.SharedPoolId = ""
		})
		resolverCalls := ts.resourcePoolResolver.calls.Load()

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:    cluster.GetClusterId(),
			TargetCmek:   newTestAwsCmekFull(),
			ValidateOnly: true,
		})
		require.NotNil(t, bizErr)
		assert.Contains(t, bizErr.Error(), "does not have shared pool id")
		assert.Equal(t, resolverCalls, ts.resourcePoolResolver.calls.Load(),
			"switching an existing cluster must not fall back to the default shared pool")
	})

	t.Run("ValidateOnly", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		// Two CheckKms hits: first from switch.go:94 verify, second from
		// CreateCustomerManagedEncryptionKey which re-verifies internally.
		expectCheckKmsSharedPoolIDTwice(t, ts, cluster.GetSharedPoolId())

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:    cluster.GetClusterId(),
			TargetCmek:   newTestAwsCmekFull(),
			ValidateOnly: true,
		})
		require.Nil(t, bizErr)
	})

	t.Run("EnableCMEK_Success", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)

		// Workflow mock: expect RunDagWithBizId to be called
		ts.mockWorkflow.EXPECT().
			RunDagWithBizId(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil)
		// Two CheckKms hits: first from switch.go:94 verify, second from
		// CreateCustomerManagedEncryptionKey which re-verifies internally.
		expectCheckKmsSharedPoolIDTwice(t, ts, cluster.GetSharedPoolId())

		cmek := newTestAwsCmekFull()

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:  cluster.GetClusterId(),
			TargetCmek: cmek,
		})
		require.Nil(t, bizErr)

		// Verify binding was created in Modifying state
		binding, bindErr := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		require.Nil(t, bindErr)
		assert.Equal(t, entity.ClusterEncryptionBindingStateModifying, binding.State)
		assert.NotZero(t, binding.EncryptionKeyId)
	})

	t.Run("EnableDefaultKey_Success", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid)

		// 2 calls: one from createDefaultKey, one from SwitchClusterEncryptionKey
		ts.mockWorkflow.EXPECT().
			RunDagWithBizId(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil).
			Times(2)

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:        cluster.GetClusterId(),
			TargetDefaultKey: true,
		})
		require.Nil(t, bizErr)

		// Verify binding was created
		binding, bindErr := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		require.Nil(t, bindErr)
		assert.Equal(t, entity.ClusterEncryptionBindingStateModifying, binding.State)
		assert.NotZero(t, binding.EncryptionKeyId)
	})

	t.Run("SwitchCMEK_to_DefaultKey", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)

		// Create cluster with an existing CMEK
		existingKey := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid)
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.EncryptionKeyId = existingKey.Id
		})

		// 2 calls: one from createDefaultKey, one from SwitchClusterEncryptionKey
		ts.mockWorkflow.EXPECT().
			RunDagWithBizId(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil).
			Times(2)

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:        cluster.GetClusterId(),
			TargetDefaultKey: true,
		})
		require.Nil(t, bizErr)

		binding, _ := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		assert.Equal(t, entity.ClusterEncryptionBindingStateModifying, binding.State)
	})

	t.Run("SwitchDefaultKey_to_CMEK", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)

		// Create cluster with a default key
		defaultKey := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.CustomerManaged = false
			k.MasterKey = "default_key"
		})
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.EncryptionKeyId = defaultKey.Id
		})

		ts.mockWorkflow.EXPECT().
			RunDagWithBizId(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil)
		// Two CheckKms hits: first from switch.go:94 verify, second from
		// CreateCustomerManagedEncryptionKey which re-verifies internally.
		expectCheckKmsSharedPoolIDTwice(t, ts, cluster.GetSharedPoolId())

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:  cluster.GetClusterId(),
			TargetCmek: newTestAwsCmekFull(),
		})
		require.Nil(t, bizErr)

		binding, _ := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		assert.Equal(t, entity.ClusterEncryptionBindingStateModifying, binding.State)
	})

	t.Run("ReplaceCMEK_A_to_B", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)

		cmekA := createTestEncryptionKey(ctx, t, ts.encryptionKeyRepo, sfid, func(k *entity.EncryptionKey) {
			k.MasterKey = "arn:aws:kms:us-west-2:123456789012:key/key-a"
		})
		cluster := createTestCluster(ctx, t, ts.clusterRepo, sfid, func(c *entity.Cluster) {
			c.EncryptionKeyId = cmekA.Id
		})

		ts.mockWorkflow.EXPECT().
			RunDagWithBizId(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil)
		// Two CheckKms hits: first from switch.go:94 verify, second from
		// CreateCustomerManagedEncryptionKey which re-verifies internally.
		expectCheckKmsSharedPoolIDTwice(t, ts, cluster.GetSharedPoolId())

		cmekB := newTestAwsCmekFull()
		cmekB.GetAwsKms().KmsKeyArn = "arn:aws:kms:us-west-2:123456789012:key/key-b"

		bizErr := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:  cluster.GetClusterId(),
			TargetCmek: cmekB,
		})
		require.Nil(t, bizErr)

		binding, _ := ts.bindingRepo.LoadByClusterId(ctx, cluster.GetClusterId())
		assert.Equal(t, entity.ClusterEncryptionBindingStateModifying, binding.State)
	})

	t.Run("MultiCluster_IndependentOperations", func(t *testing.T) {
		cleanTestData(ctx, t)
		ts := newTestService(t)

		cluster1 := createTestCluster(ctx, t, ts.clusterRepo, sfid)
		cluster2 := createTestCluster(ctx, t, ts.clusterRepo, sfid)

		// 3 calls: cluster1 CMEK=1, cluster2 TargetDefaultKey=(createDefaultKey+switch=2)
		ts.mockWorkflow.EXPECT().
			RunDagWithBizId(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
			Return(nil).
			Times(3)
		// Only cluster1 (AWS CMEK) hits CheckKms; cluster2 uses default key.
		// Two CheckKms hits: first from switch.go:94 verify, second from
		// CreateCustomerManagedEncryptionKey which re-verifies internally.
		expectCheckKmsSharedPoolIDTwice(t, ts, cluster1.GetSharedPoolId())

		// Start operation on cluster1
		err1 := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:  cluster1.GetClusterId(),
			TargetCmek: newTestAwsCmekFull(),
		})
		require.Nil(t, err1)

		// Start operation on cluster2 — should succeed independently
		err2 := ts.svc.SwitchClusterEncryptionKey(ctx, &SwitchEncryptionRequest{
			ClusterId:        cluster2.GetClusterId(),
			TargetDefaultKey: true,
		})
		require.Nil(t, err2)

		b1, _ := ts.bindingRepo.LoadByClusterId(ctx, cluster1.GetClusterId())
		b2, _ := ts.bindingRepo.LoadByClusterId(ctx, cluster2.GetClusterId())
		assert.Equal(t, cluster1.GetClusterId(), b1.ClusterId)
		assert.Equal(t, cluster2.GetClusterId(), b2.ClusterId)
	})
}

func expectCheckKmsSharedPoolIDTwice(t *testing.T, ts *testService, sharedPoolID string) {
	t.Helper()
	ts.mockRegionalCli.EXPECT().
		CheckKms(gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			_ *entity.Region,
			req *regionalserveradapter.CheckKmsRequest,
		) bizerr.BizError {
			assert.Equal(t, sharedPoolID, req.SharedPoolId)
			return nil
		}).
		Times(2)
}

// newTestAwsCmek returns a minimal CMEK proto for request building.
func newTestAwsCmek() *globalngv1.CustomerManagedEncryptionKey {
	return &globalngv1.CustomerManagedEncryptionKey{
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		RegionId:        globalcommonv1.Region_AWS_US_WEST_2,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		OrganizationId:  testOrgID,
	}
}

// newTestAwsCmekFull returns a complete CMEK proto with key, principal, and all fields.
func newTestAwsCmekFull() *globalngv1.CustomerManagedEncryptionKey {
	return &globalngv1.CustomerManagedEncryptionKey{
		CloudProviderId: globalcommonv1.CloudProvider_AWS,
		RegionId:        globalcommonv1.Region_AWS_US_WEST_2,
		ServicePlan:     globalcommonv1.ServicePlan_Premium,
		OrganizationId:  testOrgID,
		Key: &globalngv1.CustomerManagedEncryptionKey_AwsKms_{
			AwsKms: &globalngv1.CustomerManagedEncryptionKey_AwsKms{
				KmsKeyArn: "arn:aws:kms:us-west-2:123456789012:key/test-key-id",
			},
		},
		Principal: &globalngv1.CustomerManagedEncryptionKey_AwsPrincipal{
			AwsPrincipal: &globalngv1.CmekAccessIamPrincipal_AwsCmekPrincipal{
				AccountId: "123456789012",
			},
		},
	}
}