package server

import (
	"context"
	"os"
	"testing"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	gomock "go.uber.org/mock/gomock"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/yaml"

	"github.com/tidbcloud/cluster-controller/server/pkg/config"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/cloud/azure"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/scheme"
)

const testClusterID = "123456"

func TestClusterServer_CreateCluster(t *testing.T) {
	type fields struct {
		KubeClient client.Client
	}
	type args struct {
		ctx context.Context
		req *apiv1.CreateClusterRequest
	}
	tests := []struct {
		name           string
		fields         fields
		args           args
		want           *cloudv1.Cluster
		wantJWKsSecret *corev1.Secret
		wantErr        bool
	}{
		{
			name: "basic",
			fields: fields{
				KubeClient: fake.NewClientBuilder().WithScheme(scheme.Scheme).Build(),
			},
			args: args{
				ctx: nil,
				req: &apiv1.CreateClusterRequest{
					Cluster: &apiv1.Cluster{
						Id:                 "123",
						TenantId:           "456",
						ProjectId:          "789",
						DefaultTidbGroupId: defaultLoadBalancerID,
						ComponentNodeSpecs: []*apiv1.Cluster_NodeSpec{
							{
								ComponentType: apicommon.ComponentType_TIPROXY,
								InstanceType:  "c5.xlarge",
								Count:         3,
							},
							{
								ComponentType: apicommon.ComponentType_TIDB,
								InstanceType:  "c5.xlarge",
								Count:         3,
							},
							{
								ComponentType: apicommon.ComponentType_TIKV,
								InstanceType:  "r5.xlarge",
								Count:         6,
								Storage: &apiv1.Cluster_Storage{
									Type:          "gp3",
									SizeGi:        100,
									Iops:          3000,
									ThroughputMbs: 125,
								},
								AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
									{
										Name:          "raft",
										Type:          "io2",
										SizeGi:        200,
										Iops:          6000,
										ThroughputMbs: 250,
										MountPath:     "/var/lib/tikv-raft",
									},
								},
							},
							{
								ComponentType: apicommon.ComponentType_PD,
								InstanceType:  "c5.xlarge",
								Count:         3,
								Storage: &apiv1.Cluster_Storage{
									Type:          "gp3",
									SizeGi:        100,
									Iops:          3000,
									ThroughputMbs: 125,
								},
							},
						},
						ComponentConfigurations: []*apiv1.ComponentConfiguration{
							{
								ConfigurationFileContents: "security=123",
								ComponentType:             apicommon.ComponentType_TIDB,
							},
						},
						ClusterVersion: &apiv1.ClusterVersion{
							ClusterVersion: "v7.5.0",
							ComponentVersions: []*apiv1.ClusterVersion_ComponentVersion{
								{
									ComponentType: apicommon.ComponentType_TIDB,
									ImageRepo:     "pingcap/tidb",
									ImageTag:      "v7.2.1-20231127",
								},
								{
									ComponentType: apicommon.ComponentType_TIKV,
									ImageRepo:     "pingcap/tikv",
									ImageTag:      "v7.2.1-20231127",
								},
							},
						},
						SecurityConfigurations: &apiv1.Cluster_SecurityConfigurations{
							Cmek: &apiv1.Cluster_SecurityConfigurations_CmekAwsMasterKeyArn{
								CmekAwsMasterKeyArn: "cmekAwsMasterKeyArn",
							},
							TidbAuthTokenJwks: "tidbauthTokenJwks",
						},
						ClusterIdentifierHash: "clientidentifierhash",
						Port:                  4000,
						InitSqlStatements:     "select 1",
						Kubernetes: &apiv1.Kubernetes{
							Version:   "v1.27",
							Endpoint:  "https://k8s.com",
							Ca:        "ca",
							AccountId: "3804111234567",
							Vpc:       "vpc-test123",
							Info: &apiv1.Kubernetes_Aws_{
								Aws: &apiv1.Kubernetes_Aws{
									SecurityGroup: "sg-123",
									NodeSubnets: []*apiv1.Kubernetes_AwsSubnet{
										{
											Id:     "subnet-a",
											AzId:   "usw2-az1",
											AzName: "us-west-2a",
										},
										{
											Id:     "subnet-b",
											AzId:   "usw2-az3",
											AzName: "us-west-2b",
										},
										{
											Id:     "subnet-c",
											AzId:   "usw2-az4",
											AzName: "us-west-2c",
										},
									},
									PublicLoadbalancerSubnets: []*apiv1.Kubernetes_AwsSubnet{
										{
											Id:     "subnet-a",
											AzId:   "usw2-az1",
											AzName: "us-west-2a",
										},
										{
											Id:     "subnet-b",
											AzId:   "usw2-az3",
											AzName: "us-west-2b",
										},
										{
											Id:     "subnet-c",
											AzId:   "usw2-az4",
											AzName: "us-west-2c",
										},
									},
									PrivateLoadbalancerSubnets: []*apiv1.Kubernetes_AwsSubnet{
										{
											Id:     "subnet-a",
											AzId:   "usw2-az1",
											AzName: "us-west-2a",
										},
										{
											Id:     "subnet-b",
											AzId:   "usw2-az3",
											AzName: "us-west-2b",
										},
										{
											Id:     "subnet-c",
											AzId:   "usw2-az4",
											AzName: "us-west-2c",
										},
									},
									AdminNodeRole: "admin",
								},
							},
						},
					},
				},
			},
			want: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:              "123",
					Namespace:         "nightly-ms-789",
					ResourceVersion:   "1",
					CreationTimestamp: metav1.Time{},
					Labels: map[string]string{
						"cloud.pingcap.com/provider-type":         "eks",
						"cloud.pingcap.com/region":                "us-west-2",
						"cloud.pingcap.com/want-token-based-auth": "true",
						"tags.tidbcloud.com/cluster":              "123",
						"tags.tidbcloud.com/project":              "789",
						"tags.tidbcloud.com/servicetype":          "dedicated",
						"tags.tidbcloud.com/tenant":               "456",
						"creator":                                 "cluster-controller-server",
					},
					Annotations: map[string]string{
						"guard.tidbcloud.com/version":                              "v7.5.0",
						"tidbcloud.com/default-tidbgroup-access-id":                defaultLoadBalancerID,
						"eks.infra.pingcap.com/cmek-enabled":                       "true",
						"tidbcloud.com/use-cmek":                                   "true",
						"tidbcloud.com/use-3-layer-node-pool":                      "true",
						"tidbcloud.com/use-dedicated-sg":                           "true",
						"tidbcloud.com/use-enable-tiflash-graceful-shutdown":       "true",
						"tidbcloud.com/use-enable-tiflash-graceful-shutdown-sleep": "true",
						"tidbcloud.com/use-tidb-expose-oom-path":                   "true",
						"tidbcloud.com/use-tidb-expose-replayer-path":              "true",
						"tidbcloud.com/use-tiproxy-cert-layout-v1":                 "true",
						"eks.infra.pingcap.com/override-cluster-helper-image":      "",
						"eks.infra.pingcap.com/override-cluster-registry":          "",
						"eks.infra.pingcap.com/override-cluster-repoprefix":        "",
						consts.AnnoKeyLBTargetOnTiProxy(defaultLoadBalancerID):     consts.AnnoValTrue,
					},
				},
				Spec: cloudv1.ClusterSpec{
					Kubernetes: cloudv1.Kubernetes{
						Endpoint: "https://k8s.com",
						CAData:   []byte("ca"),
						Account:  "3804111234567",
						Network:  "vpc-test123",
						Provider: "eks",
						Version:  "v1.27",
						Cloud: cloudv1.Cloud{
							AWS: &cloudv1.CloudAWS{
								EKSSecurityGroup: "sg-123",
								SubnetGroups: []cloudv1.AWSSubnetGroup{
									{
										Name: "nodes",
										Subnets: []cloudv1.AWSSubnet{
											{
												SubnetID: "subnet-a",
												Zone:     "us-west-2a",
											},
											{
												SubnetID: "subnet-b",
												Zone:     "us-west-2b",
											},
											{
												SubnetID: "subnet-c",
												Zone:     "us-west-2c",
											},
										},
									},
									{
										Name: "public-nlb",
										Subnets: []cloudv1.AWSSubnet{
											{
												SubnetID: "subnet-a",
												Zone:     "us-west-2a",
											},
											{
												SubnetID: "subnet-b",
												Zone:     "us-west-2b",
											},
											{
												SubnetID: "subnet-c",
												Zone:     "us-west-2c",
											},
										},
									},
									{
										Name: "private-nlb",
										Subnets: []cloudv1.AWSSubnet{
											{
												SubnetID: "subnet-a",
												Zone:     "us-west-2a",
											},
											{
												SubnetID: "subnet-b",
												Zone:     "us-west-2b",
											},
											{
												SubnetID: "subnet-c",
												Zone:     "us-west-2c",
											},
										},
									},
								},
								AdminNodeRole: "admin",
							},
						},
					},
					SqlEndpoints: cloudv1.SqlEndpoints{
						DNS:                  cloudv1.DNS{Identity: "clientidentifierhash"},
						Public:               cloudv1.Service{Enabled: ptr.To[bool](false)},
						Private:              &cloudv1.Service{Enabled: ptr.To[bool](false)},
						UseTiProxyAsLBTarget: true,
					},
					Cluster: cloudv1.ClusterTopology{
						TiProxy: cloudv1.TiProxy{
							Replicas:     3,
							InstanceType: "c5.xlarge",
						},
						TiDB: cloudv1.TiDB{
							Replicas:     3,
							InstanceType: "c5.xlarge",
							Security: &cloudv1.TiDBSecurity{
								AuthTokenJwksSecret: &corev1.SecretReference{
									Name:      "123-tidb-auth-token-jwks-secret",
									Namespace: "nightly-ms-789",
								},
							},
							Audit: &cloudv1.Audit{
								Enable: true,
							},
						},
						PD: cloudv1.PD{
							Replicas:         3,
							InstanceType:     "c5.xlarge",
							StorageClassName: "pd-gp3-3000-125",
							StorageSize:      storageSizeGi2quantity(100),
						},
						TiKV: cloudv1.TiKV{
							Replicas:         6,
							InstanceType:     "r5.xlarge",
							StorageClassName: "tikv-gp3-3000-125",
							StorageSize:      storageSizeGi2quantity(100),
							StorageVolumes: []cloudv1.ComponentStorageVolume{
								{
									Name:             "raft",
									StorageClassName: "tikv-io2-6000-250",
									StorageSize:      storageSizeGi2quantity(200),
									MountPath:        "/var/lib/tikv-raft",
								},
							},
						},
						TiFlash: cloudv1.TiFlash{
							StorageSize: resource.Quantity{
								Format: resource.DecimalSI,
							},
						},
					},
					RecoveryMode: false,
					Paused:       nil,
					ForProvider: &runtime.RawExtension{
						Raw: []byte(`{"apiVersion":"eks.infra.pingcap.com/v1alpha1","backupStorage":{},"kind":"EKSCluster","restoreStorage":{},"security":{"masterKeyArn":"cmekAwsMasterKeyArn"}}`),
					},
					ConfigRevisionRef: &cloudv1.Reference{
						Name: "cr--t456-c123-f013913c",
					},
				},
			},
			wantErr: false,
		},
	}

	ns := "nightly-ms-789"
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &ClusterServer{
				KubeClient: tt.fields.KubeClient,
			}
			_, err := s.CreateCluster(tt.args.ctx, tt.args.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("ClusterServer.CreateCluster() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			got := &cloudv1.Cluster{}
			if err := s.KubeClient.Get(tt.args.ctx, client.ObjectKey{Namespace: ns, Name: tt.args.req.Cluster.Id}, got); err != nil {
				t.Errorf("kubeclient get error = %v", err)
				return
			}

			// ignore status field
			got.Status = cloudv1.ClusterStatus{}
			// ignore configRevisionRef due to configRevision.Name contains hash of timestamp
			got.Spec.ConfigRevisionRef = tt.want.Spec.ConfigRevisionRef

			gotYamlData, err := yaml.Marshal(got)
			if err != nil {
				t.Logf("Error: %v", err)
				return
			}
			// t.Errorf(string(gotYamlData))

			wantYamlData, err := yaml.Marshal(tt.want)
			if err != nil {
				t.Logf("Error: %v", err)
				return
			}
			assert.Equalf(t, string(wantYamlData), string(gotYamlData), "cluster CR")

			if tt.wantJWKsSecret != nil {
				jwksSecret := &corev1.Secret{}
				if err := s.KubeClient.Get(tt.args.ctx, client.ObjectKey{Namespace: ns, Name: got.GetJWKsSecretName()}, jwksSecret); err != nil {
					t.Errorf("kubeclient get error = %v", err)
					return
				}
				assert.Equalf(t, jwksSecret.Data, map[string][]byte{
					cloudv1.TiDBAuthTokenJWKSFilename: []byte(tt.args.req.Cluster.SecurityConfigurations.TidbAuthTokenJwks),
				}, "jwks secret")
			}
		})
	}
}

func TestMain(m *testing.M) {
	os.Setenv("PROXY_AUTH0_CLIENT_ID", "test")
	os.Setenv("PROXY_AUTH0_CLIENT_SECRET", "test")
	if err := config.InitConfig(nil); err != nil {
		panic(err)
	}

	code := m.Run()
	os.Exit(code)
}

func TestClusterServer_validateSetAuditLogConfiguration(t *testing.T) {
	type args struct {
		req     func(*apiv1.SetAuditLogConfigurationRequest)
		cluster func(*cloudv1.Cluster)
	}
	tests := []struct {
		name      string
		args      args
		errAssert func(t assert.TestingT, err error)
	}{
		{
			name: "should raise error if AuditLogConfiguration is nil when the bucket is managed by customer",
			args: args{
				req: func(request *apiv1.SetAuditLogConfigurationRequest) {
					request.ManagedBy = apiv1.SetAuditLogConfigurationRequest_BUCKET_MANAGED_BY_CUSTOMER
					request.AuditLogConfiguration = nil
				},
			},
			errAssert: func(t assert.TestingT, err error) {
				assert.ErrorContains(t, err, "AuditLogConfiguration is required")
			},
		},
		{
			name: "should raise error if shared bucket is not enabled but required to use PingCAP managed bucket",
			args: args{
				req: func(request *apiv1.SetAuditLogConfigurationRequest) {
					request.ManagedBy = apiv1.SetAuditLogConfigurationRequest_BUCKET_MANAGED_BY_PINGCAP
				},
				cluster: func(cluster *cloudv1.Cluster) {
					cluster.Spec.SharedBucket = nil
				},
			},
			errAssert: func(t assert.TestingT, err error) {
				assert.ErrorContains(t, err, "shared bucket is not enabled")
			},
		},
		{
			name: "should raise error if shared bucket is not ready but required to use PingCAP managed bucket",
			args: args{
				req: func(request *apiv1.SetAuditLogConfigurationRequest) {
					request.ManagedBy = apiv1.SetAuditLogConfigurationRequest_BUCKET_MANAGED_BY_PINGCAP
				},
				cluster: func(cluster *cloudv1.Cluster) {
					cluster.Status.SharedBucket.AWS.DBAuditLogWriterRoleARN = ""
				},
			},
			errAssert: func(t assert.TestingT, err error) {
				assert.ErrorContains(t, err, "shared bucket is not ready")
			},
		},
		{
			name: "normal case for PingCAP managed bucket",
			args: args{
				req: func(request *apiv1.SetAuditLogConfigurationRequest) {
					request.ManagedBy = apiv1.SetAuditLogConfigurationRequest_BUCKET_MANAGED_BY_PINGCAP
				},
			},
			errAssert: func(t assert.TestingT, err error) {
				assert.NoError(t, err)
			},
		},
		{
			name: "normal case for Customer managed bucket",
			args: args{
				req: func(request *apiv1.SetAuditLogConfigurationRequest) {
					request.ManagedBy = apiv1.SetAuditLogConfigurationRequest_BUCKET_MANAGED_BY_CUSTOMER
				},
				cluster: func(cluster *cloudv1.Cluster) {
					cluster.Spec.SharedBucket = nil
				},
			},
			errAssert: func(t assert.TestingT, err error) {
				assert.NoError(t, err)
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := &apiv1.SetAuditLogConfigurationRequest{
				AuditLogConfiguration: &apiv1.AuditLogConfiguration{
					BucketUrl:    "s3:bucket/t1/p1",
					BucketRegion: "us-west-2",
					Storage: &apiv1.AuditLogConfiguration_S3_{
						S3: &apiv1.AuditLogConfiguration_S3{
							RoleArn: "role1",
						},
					},
				},
				ManagedBy: apiv1.SetAuditLogConfigurationRequest_BUCKET_MANAGED_BY_CUSTOMER,
			}
			cluster := &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{
						cloudv1.TenantTagKey:      "t1",
						cloudv1.ProjectTagKey:     "p1",
						cloudv1.ProviderTypeLabel: cloudv1.ProviderTypeEKS,
					},
					Name: "c1",
				},
				Spec: cloudv1.ClusterSpec{
					SharedBucket: &cloudv1.SharedBucket{
						BucketURL: "s3:sb/t1/p1/c1",
					},
				},
				Status: cloudv1.ClusterStatus{
					SharedBucket: &cloudv1.SharedBucketStatus{
						AWS: &cloudv1.AWSSharedBucketStatus{
							DBAuditLogWriterRoleARN: "role2",
						},
					},
				},
			}
			if tt.args.cluster != nil {
				tt.args.cluster(cluster)
			}
			if tt.args.req != nil {
				tt.args.req(req)
			}
			s := &ClusterServer{}
			tt.errAssert(t, s.validateSetAuditLogConfiguration(req, cluster))
		})
	}
}

func TestClusterServer_GetControlplaneK8SResource(t *testing.T) {
	type args struct {
		ctx      context.Context
		queryReq *apiv1.GetK8SResourceRequest
	}
	gvk_1 := schema.GroupVersionKind{
		Group:   "cloud.pingcap.com",
		Version: "v1",
		Kind:    "Cluster",
	}
	gkv_2 := schema.GroupVersionKind{
		Group:   "",
		Version: "v1",
		Kind:    "Service",
	}
	mapper := meta.NewDefaultRESTMapper([]schema.GroupVersion{gvk_1.GroupVersion(), gkv_2.GroupVersion()})
	mapper.Add(gvk_1, meta.RESTScopeNamespace)
	mapper.Add(gkv_2, meta.RESTScopeNamespace)
	kubeClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRESTMapper(mapper).Build()
	testCluster := &apiv1.CreateClusterRequest{
		Cluster: &apiv1.Cluster{
			Id:        "123",
			TenantId:  "456",
			ProjectId: "789",
			ComponentNodeSpecs: []*apiv1.Cluster_NodeSpec{
				{
					ComponentType: apicommon.ComponentType_TIPROXY,
					InstanceType:  "c5.xlarge",
					Count:         3,
				},
				{
					ComponentType: apicommon.ComponentType_TIDB,
					InstanceType:  "c5.xlarge",
					Count:         3,
				},
				{
					ComponentType: apicommon.ComponentType_TIKV,
					InstanceType:  "r5.xlarge",
					Count:         6,
					Storage: &apiv1.Cluster_Storage{
						Type:          "gp3",
						SizeGi:        100,
						Iops:          3000,
						ThroughputMbs: 125,
					},
					AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
						{
							Name:          "raft",
							Type:          "io2",
							SizeGi:        200,
							Iops:          6000,
							ThroughputMbs: 250,
							MountPath:     "/var/lib/tikv-raft",
						},
					},
				},
				{
					ComponentType: apicommon.ComponentType_PD,
					InstanceType:  "c5.xlarge",
					Count:         3,
					Storage: &apiv1.Cluster_Storage{
						Type:          "gp3",
						SizeGi:        100,
						Iops:          3000,
						ThroughputMbs: 125,
					},
				},
			},
			ComponentConfigurations: []*apiv1.ComponentConfiguration{
				{
					ConfigurationFileContents: "security=123",
					ComponentType:             apicommon.ComponentType_TIDB,
				},
			},
			ClusterVersion: &apiv1.ClusterVersion{
				ClusterVersion: "v7.5.0",
				ComponentVersions: []*apiv1.ClusterVersion_ComponentVersion{
					{
						ComponentType: apicommon.ComponentType_TIDB,
						ImageRepo:     "pingcap/tidb",
						ImageTag:      "v7.2.1-20231127",
					},
					{
						ComponentType: apicommon.ComponentType_TIKV,
						ImageRepo:     "pingcap/tikv",
						ImageTag:      "v7.2.1-20231127",
					},
				},
			},
			SecurityConfigurations: &apiv1.Cluster_SecurityConfigurations{
				Cmek: &apiv1.Cluster_SecurityConfigurations_CmekAwsMasterKeyArn{
					CmekAwsMasterKeyArn: "cmekAwsMasterKeyArn",
				},
				TidbAuthTokenJwks: "tidbauthTokenJwks",
			},
			ClusterIdentifierHash: "clientidentifierhash",
			Port:                  4000,
			InitSqlStatements:     "select 1",
			Kubernetes: &apiv1.Kubernetes{
				Version:   "v1.27",
				Endpoint:  "https://k8s.com",
				Ca:        "ca",
				AccountId: "3804111234567",
				Vpc:       "vpc-test123",
				Info: &apiv1.Kubernetes_Aws_{
					Aws: &apiv1.Kubernetes_Aws{
						SecurityGroup: "sg-123",
						NodeSubnets: []*apiv1.Kubernetes_AwsSubnet{
							{
								Id:     "subnet-a",
								AzId:   "usw2-az1",
								AzName: "us-west-2a",
							},
							{
								Id:     "subnet-b",
								AzId:   "usw2-az3",
								AzName: "us-west-2b",
							},
							{
								Id:     "subnet-c",
								AzId:   "usw2-az4",
								AzName: "us-west-2c",
							},
						},
						PublicLoadbalancerSubnets: []*apiv1.Kubernetes_AwsSubnet{
							{
								Id:     "subnet-a",
								AzId:   "usw2-az1",
								AzName: "us-west-2a",
							},
							{
								Id:     "subnet-b",
								AzId:   "usw2-az3",
								AzName: "us-west-2b",
							},
							{
								Id:     "subnet-c",
								AzId:   "usw2-az4",
								AzName: "us-west-2c",
							},
						},
						PrivateLoadbalancerSubnets: []*apiv1.Kubernetes_AwsSubnet{
							{
								Id:     "subnet-a",
								AzId:   "usw2-az1",
								AzName: "us-west-2a",
							},
							{
								Id:     "subnet-b",
								AzId:   "usw2-az3",
								AzName: "us-west-2b",
							},
							{
								Id:     "subnet-c",
								AzId:   "usw2-az4",
								AzName: "us-west-2c",
							},
						},
						AdminNodeRole: "admin",
					},
				},
			},
		},
	}
	testService := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "test-svc",
			Namespace: "default",
			Labels: map[string]string{
				cloudv1.ClusterTagKey: "test-id",
			},
		},
		Spec: corev1.ServiceSpec{
			Ports: []corev1.ServicePort{
				{Port: 80},
			},
		},
	}
	s := &ClusterServer{
		KubeClient: kubeClient,
	}
	_, err := s.CreateCluster(context.TODO(), testCluster)
	if err != nil {
		t.Fatalf("Failed to create test cluster: %v", err)
	}
	err = s.KubeClient.Create(context.TODO(), testService)
	if err != nil {
		t.Fatalf("Failed to create test service: %v", err)
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		{
			name: "basic",
			args: args{
				ctx: context.TODO(),
				queryReq: &apiv1.GetK8SResourceRequest{
					ClusterId:    "123",
					Namespace:    "nightly-ms-789",
					ResourceKind: "Cluster",
					ResourceName: "123",
				},
			},
			wantErr: false,
		},
		{
			name: "not found",
			args: args{
				ctx: context.TODO(),
				queryReq: &apiv1.GetK8SResourceRequest{
					ClusterId:    "123",
					Namespace:    "nightly-ms-789",
					ResourceKind: "Cluster",
					ResourceName: "321",
				},
			},
			wantErr: true,
		},
		{
			name: "k8s native component",
			args: args{
				ctx: context.TODO(),
				queryReq: &apiv1.GetK8SResourceRequest{
					ClusterId:    "test-id",
					Namespace:    "default",
					ResourceKind: "Service",
					ResourceName: "test-svc",
				},
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := s.GetControlplaneK8SResource(tt.args.ctx, tt.args.queryReq)
			if (err != nil) != tt.wantErr {
				t.Errorf("ClusterServer.GetControlplaneMeta() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if tt.wantErr {
				t.Logf("ClusterServer.GetControlplaneMeta() got error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got == nil {
				t.Errorf("ClusterServer.GetControlplaneMeta() got = nil, want non-nil")
				return
			}
			gotYamlData, err := yaml.Marshal(got)
			if err != nil {
				t.Logf("Error: %v", err)
				return
			}
			t.Logf("Got:\n%s", string(gotYamlData))
		})
	}
}

func TestClusterServer_buildCMEKForAzure(t *testing.T) {
	type args struct {
		ctx     context.Context
		cluster *cloudv1.Cluster
		conf    *apiv1.Cluster_SecurityConfigurations
	}
	tests := []struct {
		name            string
		args            args
		wantClusterSpec *cloudv1.AKSClusterSpec
		wantChanged     bool
		wantErr         bool
		keyVersions     []*azkeys.KeyProperties
		beforeFunc      func(t *testing.T, args *args)
	}{
		{
			name: "azure cmek is nil, no change",
			args: args{
				cluster: &cloudv1.Cluster{},
				conf:    &apiv1.Cluster_SecurityConfigurations{},
			},
			wantChanged: false,
			wantErr:     false,
		},
		{
			name: "empty keyvault name, raise error",
			args: args{
				cluster: &cloudv1.Cluster{},
				conf: &apiv1.Cluster_SecurityConfigurations{
					Cmek: &apiv1.Cluster_SecurityConfigurations_AzureCmek_{
						AzureCmek: &apiv1.Cluster_SecurityConfigurations_AzureCmek{
							ApplicationId:             "",
							ManagedIdentityResourceId: "",
							KeyVaultName:              "",
							KeyName:                   "mykey",
						},
					},
				},
			},
			wantChanged: false,
			wantErr:     true,
		},
		{
			name: "valid keyvault and key, query latest version",
			args: args{
				cluster: &cloudv1.Cluster{},
				conf: &apiv1.Cluster_SecurityConfigurations{
					Cmek: &apiv1.Cluster_SecurityConfigurations_AzureCmek_{
						AzureCmek: &apiv1.Cluster_SecurityConfigurations_AzureCmek{
							ApplicationId:             "test-app-id",
							ManagedIdentityResourceId: "test-resource-id",
							KeyVaultName:              "test",
							KeyName:                   "test-key",
						},
					},
				},
			},
			wantChanged: true,
			wantClusterSpec: &cloudv1.AKSClusterSpec{
				TypeMeta: metav1.TypeMeta{
					Kind:       "AKSCluster",
					APIVersion: "aks.infra.pingcap.com/v1alpha1",
				},
				AzureCMEK: &cloudv1.AzureCMEKSpec{
					ApplicationID:             "test-app-id",
					ManagedIdentityResourceID: "test-resource-id",
					KeyVaultName:              "test",
					KeyName:                   "test-key",
					KeyVersion:                "latest-version",
				},
			},
			wantErr: false,
			keyVersions: []*azkeys.KeyProperties{
				{
					KID: lo.ToPtr(azkeys.ID("https://test.vault.azure.net/keys/test-key/old-key")),
					Attributes: &azkeys.KeyAttributes{
						Created: ptr.To(time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)),
						Enabled: ptr.To(true),
					},
				},
				{
					KID: lo.ToPtr(azkeys.ID("https://test.vault.azure.net/keys/test-key/new-key-but-disabled")),
					Attributes: &azkeys.KeyAttributes{
						Created: ptr.To(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)),
						Enabled: ptr.To(false),
					},
				},
				// picked
				{
					KID: lo.ToPtr(azkeys.ID("https://test.vault.azure.net/keys/test-key/latest-version")),
					Attributes: &azkeys.KeyAttributes{
						Created: ptr.To(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)),
						Enabled: ptr.To(true),
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			s := &ClusterServer{}
			if tt.beforeFunc != nil {
				tt.beforeFunc(t, &tt.args)
			}

			if tt.keyVersions != nil {
				// make azidentity.NewWorkloadIdentityCredential happy
				os.Setenv("AZURE_CLIENT_ID", "test-client-id")
				os.Setenv("AZURE_FEDERATED_TOKEN_FILE", "/tmp/azure_token")
				os.Setenv("AZURE_TENANT_ID", "test-tenant-id")

				mockAzureCli := azure.NewMockClient(ctrl)
				s.azureCli = mockAzureCli
				mockAzureCli.EXPECT().WithCredential(gomock.Any()).Return(mockAzureCli)
				mockAzureCli.EXPECT().
					QueryKeyvaultKeyVersions(gomock.Any(), gomock.Any(), gomock.Any()).
					Return(tt.keyVersions, nil)
			}

			gotChanged, err := s.buildCMEKForAzure(tt.args.ctx, tt.args.cluster, tt.args.conf)
			if (err != nil) != tt.wantErr {
				t.Errorf("ClusterServer.buildCMEKForAzure() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			assert.Equal(t, tt.wantChanged, gotChanged, "changed")

			if tt.wantClusterSpec != nil {
				aksClusterSpec, err := cloudv1.AKSClusterSpecForCluster(tt.args.cluster)
				assert.NoError(t, err)
				assert.Equal(t, tt.wantClusterSpec, aksClusterSpec)
			}
		})
	}
}

const (
	tAzureVaultName  = "myvault"
	tAzureKeyName    = "mykey"
	tAzureKeyVersion = "123456"
)

func TestClusterServer_validateAzureCMEK(t *testing.T) {
	type args struct {
		ctx     context.Context
		cmekReq *apiv1.Cluster_SecurityConfigurations_AzureCmek
	}
	tests := []struct {
		name        string
		args        args
		wantValid   bool
		wantMessage string
		wantErr     bool
		mockAzure   func(mockAzureCli *azure.MockClient)
	}{
		{
			name: "valid keyvault and key, query latest version, valid permission",
			args: args{
				cmekReq: &apiv1.Cluster_SecurityConfigurations_AzureCmek{
					ApplicationId:             "test-app-id",
					ManagedIdentityResourceId: "test-resource-id",
					KeyVaultName:              "myvault",
					KeyName:                   "mykey",
				},
			},
			wantValid:   true,
			wantMessage: "OK",
			wantErr:     false,
			mockAzure: func(mockAzureCli *azure.MockClient) {
				keys := []*azkeys.KeyProperties{
					{
						KID: lo.ToPtr(azkeys.ID("https://myvault.vault.azure.net/keys/mykey/123456")),
						Attributes: &azkeys.KeyAttributes{
							Created: ptr.To(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)),
							Enabled: ptr.To(true),
						},
					},
				}
				wrapOperation := azkeys.KeyOperationParameters{
					Algorithm: lo.ToPtr(azkeys.EncryptionAlgorithmRSA15),
					Value:     []byte(dummyAzureEncryptionKey),
				}
				wrappedResult := []byte("abcdefg")
				wraped := &azkeys.WrapKeyResponse{
					KeyOperationResult: azkeys.KeyOperationResult{
						Result: wrappedResult,
					},
				}
				unwrapOperation := azkeys.KeyOperationParameters{
					Algorithm: lo.ToPtr(azkeys.EncryptionAlgorithmRSA15),
					Value:     wrappedResult,
				}
				unwraped := &azkeys.UnwrapKeyResponse{
					KeyOperationResult: azkeys.KeyOperationResult{
						Result: []byte(dummyAzureEncryptionKey),
					},
				}
				mockAzureCli.EXPECT().QueryKeyvaultKeyVersions(gomock.Any(), tAzureVaultName, tAzureKeyName).Return(keys, nil)
				mockAzureCli.EXPECT().WrapKeyByKeyvaultKey(gomock.Any(), tAzureVaultName, tAzureKeyName, tAzureKeyVersion, wrapOperation).Return(wraped, nil)
				mockAzureCli.EXPECT().UnwrapKeyByKeyvaultKey(gomock.Any(), tAzureVaultName, tAzureKeyName, tAzureKeyVersion, unwrapOperation).Return(unwraped, nil)
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			s := &ClusterServer{}
			mockAzureCli := azure.NewMockClient(ctrl)
			s.azureCli = mockAzureCli
			// make azidentity.NewWorkloadIdentityCredential happy
			os.Setenv("AZURE_CLIENT_ID", "test-client-id")
			os.Setenv("AZURE_FEDERATED_TOKEN_FILE", "/tmp/azure_token")
			os.Setenv("AZURE_TENANT_ID", "test-tenant-id")

			mockAzureCli.EXPECT().WithCredential(gomock.Any()).Return(mockAzureCli).AnyTimes()
			tt.mockAzure(mockAzureCli)

			gotValid, gotMessage, err := s.validateAzureCMEK(tt.args.ctx, tt.args.cmekReq)
			if (err != nil) != tt.wantErr {
				t.Errorf("ClusterServer.validateAzureCMEK() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			assert.Equal(t, tt.wantValid, gotValid, "Valid")
			assert.Equal(t, tt.wantMessage, gotMessage, "Message")
		})
	}
}