package server

import (
	"context"
	"errors"
	"reflect"
	"testing"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"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"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/tiinf/monitor/tilog"

	internalutil "github.com/tidbcloud/cluster-controller/server/internal/util"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	"github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/scheme"
)

type Option func(*cloudv1.Cluster) *cloudv1.Cluster

func withClusterReplica(replica int) Option {
	return func(c *cloudv1.Cluster) *cloudv1.Cluster {
		for _, g := range c.Spec.ExtraGroups {
			g.TiDB.Replicas = replica
		}
		return c
	}
}

func WithTiproxy(groupName string) Option {
	tiproxy := cloudv1.TiProxy{
		InstanceType: "c5.2xlarge",
		Replicas:     2,
	}
	tiproxyStatus := cloudv1.TiProxyStatus{
		UpdatedReplicas: 2,
		Members: []cloudv1.TiProxyMember{
			{
				InstanceType: "c5.2xlarge",
			},
		},
	}
	return func(c *cloudv1.Cluster) *cloudv1.Cluster {
		if isDefaultLoadBalancer(groupName, c) {
			c.Spec.Cluster.TiProxy = tiproxy
			c.Status.DB.TiProxy = &tiproxyStatus
			return c
		}
		for i, group := range c.Spec.ExtraGroups {
			if group.Name != groupName {
				continue
			}

			c.Spec.ExtraGroups[i].TiProxy = tiproxy
			c.Status.ExtraEndpoints[i].TiProxy = &tiproxyStatus
		}
		return c
	}
}

func newCluster(defaultGroupID string, fn ...Option) *cloudv1.Cluster {

	cluster := &cloudv1.Cluster{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Cluster",
			APIVersion: "cloud.pingcap.com/v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:              "456",
			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",
				cloudv1.ClusterTagKey:             "456",
			},
			Annotations: map[string]string{
				"tidbcloud.com/default-tidbgroup-access-id": defaultGroupID,
			},
		},
		Spec: cloudv1.ClusterSpec{
			SqlEndpoints: cloudv1.SqlEndpoints{
				DNS: cloudv1.DNS{Identity: "clientidentifierhash"},
				Public: cloudv1.Service{
					Enabled:                  ptr.To[bool](true),
					LoadBalancerSourceRanges: []string{"0.0.0.0", "0.0.0.1"},
					Port:                     ptr.To[int32](4000),
				},
				Private: &cloudv1.Service{
					Enabled: ptr.To[bool](true),
					Port:    ptr.To[int32](4000),
				},
			},
			ExtraGroups: []cloudv1.ExtraGroupSpec{
				{
					Name: "basic",
					TiDB: cloudv1.TiDB{
						Replicas: 3,
					},
					SqlEndpoints: cloudv1.SqlEndpoints{
						DNS: cloudv1.DNS{Identity: "clientidentifierhash2"},
						Public: cloudv1.Service{
							Enabled:                  ptr.To[bool](true),
							LoadBalancerSourceRanges: []string{"192.168.0.1", "192.168.0.1"},
							Port:                     ptr.To[int32](4000),
						},
						Private: &cloudv1.Service{
							Enabled: ptr.To[bool](true),
							Port:    ptr.To[int32](4000),
						},
					},
				},
			},
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     3,
					InstanceType: "c5.xlarge",
				},
			},
		},
		Status: cloudv1.ClusterStatus{
			Phase: cloudv1.ClusterPhaseRunning,
			SqlEndpoints: cloudv1.SqlEndpointsStatus{
				PublicEndpoint: &cloudv1.SqlEndpointStatus{
					Host:       "public",
					Port:       4000,
					ResourceID: "public-resource-id",
				},
				PrivateEndpoint: &cloudv1.SqlEndpointStatus{
					Host:       "private",
					Port:       4000,
					ResourceID: "private-resource-id",
				},
			},
			ExtraEndpoints: []cloudv1.ExtraGroupStatus{
				{
					Name: "basic",
					SqlEndpointsStatus: cloudv1.SqlEndpointsStatus{
						PublicEndpoint: &cloudv1.SqlEndpointStatus{
							Host:       "public-extra-group",
							Port:       4000,
							ResourceID: "public-extra-group-id",
						},
						PrivateEndpoint: &cloudv1.SqlEndpointStatus{
							Host:       "private-extra-group",
							Port:       4000,
							ResourceID: "private-extra-group-id",
						},
					},
					TiDB: &cloudv1.TiDBStatus{
						UpdatedReplicas: 3,
						Members: []cloudv1.TiDBMember{
							{
								InstanceType: "c5.xlarge",
							},
						},
					},
				},
			},
			DB: &cloudv1.DBStatus{
				TiDB: &cloudv1.TiDBStatus{
					UpdatedReplicas: 3,
					Members: []cloudv1.TiDBMember{
						{
							InstanceType: "c5.xlarge",
						},
					},
				},
			},
		},
	}
	for _, f := range fn {
		cluster = f(cluster)
	}
	return cluster
}

func TestTidbGroupServer_CreateTidbGroup(t *testing.T) {
	type fields struct {
		KubeClient    client.Client
		clusterReader internalutil.ClusterReader
	}
	type args struct {
		ctx context.Context
		req *apiv1.CreateTidbGroupRequest
	}
	tests := []struct {
		name            string
		fields          fields
		args            args
		want            *apiv1.TidbGroup
		wantErrContains string
	}{
		{
			name: "create a new tidbgroup should be OK",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.CreateTidbGroupRequest{
					TidbGroup: &apiv1.TidbGroup{
						Id:        "123",
						ClusterId: "456",
						Tidb: &apiv1.Cluster_NodeSpec{
							InstanceType: "c5.xlarge",
							Count:        3,
						},
					},
				},
			},
			wantErrContains: "",
			want: &apiv1.TidbGroup{
				Id:        "123",
				ClusterId: "456",
				Tidb: &apiv1.Cluster_NodeSpec{
					InstanceType: "c5.xlarge",
					Count:        3,
				},
			},
		},
		{
			name: "create a new tidbgroup with backward compatibility should be OK",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.CreateTidbGroupRequest{
					TidbGroup: &apiv1.TidbGroup{
						Id:        "123",
						ClusterId: "456",
						Count:     3,
					},
				},
			},
			wantErrContains: "",
			want: &apiv1.TidbGroup{
				Id:        "123",
				ClusterId: "456",
				Count:     3,
				Tidb: &apiv1.Cluster_NodeSpec{
					InstanceType: "c5.xlarge",
					Count:        3,
				},
			},
		},
		{
			name: "create a new tidbgroup with default tidb group should be OK",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.CreateTidbGroupRequest{
					TidbGroup: &apiv1.TidbGroup{
						Id:        "123",
						ClusterId: "456",
						Tidb: &apiv1.Cluster_NodeSpec{
							InstanceType: "c5.xlarge",
							Count:        3,
						},
					},
				},
			},
			wantErrContains: "",
			want: &apiv1.TidbGroup{
				Id:        "123",
				ClusterId: "456",
				Tidb: &apiv1.Cluster_NodeSpec{
					InstanceType: "c5.xlarge",
					Count:        3,
				},
				IsDefault: true,
			},
		},
		{
			name: "miss tidb count", // this is a expected behavior that we shouldn't set the tidb count in tidbgroup creation.
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.CreateTidbGroupRequest{
					TidbGroup: &apiv1.TidbGroup{
						Id:        "group-new",
						ClusterId: "456",
						Tidb: &apiv1.Cluster_NodeSpec{
							InstanceType: "", // will use the default instance type as tidbgroup instance type
							Count:        0,
						},
					},
				},
			},
			wantErrContains: "",
			want: &apiv1.TidbGroup{
				Id:        "group-new",
				ClusterId: "456",
				Tidb: &apiv1.Cluster_NodeSpec{
					InstanceType: "c5.xlarge",
					Count:        0,
				},
				IsDefault: false,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &TidbGroupServer{
				KubeClient:    tt.fields.KubeClient,
				clusterReader: tt.fields.clusterReader,
			}
			got, err := s.CreateTidbGroup(tt.args.ctx, tt.args.req)
			if tt.wantErrContains != "" {
				assert.Contains(t, err.Error(), tt.wantErrContains)
			} else {
				assert.NoError(t, err)
			}
			assert.Equalf(t, tt.want, got, "TidbGroupServer.CreateTidbGroup() = %v, want %v", got, tt.want)
		})
	}
}

func TestTidbGroupServer_GetTidbGroup(t *testing.T) {
	type fields struct {
		KubeClient    client.Client
		clusterReader internalutil.ClusterReader
	}
	type args struct {
		ctx context.Context
		req *apiv1.GetTidbGroupRequest
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantTg  *apiv1.TidbGroup
		wantErr bool
	}{
		{
			name: "basic",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.GetTidbGroupRequest{
					TidbGroupId: "basic",
					ClusterId:   "456",
				},
			},
			wantErr: false,
			wantTg: &apiv1.TidbGroup{
				Id:                    "basic",
				ClusterId:             "456",
				ClusterIdentifierHash: "clientidentifierhash2",
				EnablePublicLb:        true,
				EnablePrivateLb:       true,
				PublicLb: &apiv1.LoadBalancer{
					Host:            "public-extra-group",
					Port:            4000,
					CloudResourceId: "public-extra-group-id",
					IpAllowList:     sourceRangesToIpAllowList([]string{"192.168.0.1", "192.168.0.1"}),
					Synced:          ptr.To(true),
				},
				PrivateLb: &apiv1.LoadBalancer{
					Host:            "private-extra-group",
					Port:            4000,
					CloudResourceId: "private-extra-group-id",
					Synced:          ptr.To(true),
				},
				InstanceType: "c5.xlarge",
				Count:        3,
				Tidb: &apiv1.Cluster_NodeSpec{
					ComponentType: common.ComponentType_TIDB,
					InstanceType:  "c5.xlarge",
					Count:         3,
				},
			},
		},
		{
			name: "default tidb group",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.GetTidbGroupRequest{
					TidbGroupId: "123",
					ClusterId:   "456",
				},
			},
			wantErr: false,
			wantTg: &apiv1.TidbGroup{
				Id:                    "123",
				ClusterId:             "456",
				IsDefault:             true,
				ClusterIdentifierHash: "clientidentifierhash",
				EnablePublicLb:        true,
				EnablePrivateLb:       true,
				PublicLb: &apiv1.LoadBalancer{
					Host:            "public",
					Port:            4000,
					CloudResourceId: "public-resource-id",
					IpAllowList:     sourceRangesToIpAllowList([]string{"0.0.0.0", "0.0.0.1"}),
					Synced:          nil,
				},
				PrivateLb: &apiv1.LoadBalancer{
					Host:            "private",
					Port:            4000,
					CloudResourceId: "private-resource-id",
					Synced:          nil,
				},
				InstanceType: "c5.xlarge",
				Count:        3,
				Tidb: &apiv1.Cluster_NodeSpec{
					ComponentType: common.ComponentType_TIDB,
					InstanceType:  "c5.xlarge",
					Count:         3,
				},
			},
		},
		{
			name: "with tiproxy in default group",
			fields: fields{
				KubeClient: fake.NewClientBuilder().WithObjects(
					newCluster("123", WithTiproxy("123")),
				).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.GetTidbGroupRequest{
					TidbGroupId: "123",
					ClusterId:   "456",
				},
			},
			wantErr: false,
			wantTg: &apiv1.TidbGroup{
				Id:                    "123",
				ClusterId:             "456",
				IsDefault:             true,
				ClusterIdentifierHash: "clientidentifierhash",
				EnablePublicLb:        true,
				EnablePrivateLb:       true,
				PublicLb: &apiv1.LoadBalancer{
					Host:            "public",
					Port:            4000,
					CloudResourceId: "public-resource-id",
					IpAllowList:     sourceRangesToIpAllowList([]string{"0.0.0.0", "0.0.0.1"}),
					Synced:          nil,
				},
				PrivateLb: &apiv1.LoadBalancer{
					Host:            "private",
					Port:            4000,
					CloudResourceId: "private-resource-id",
					Synced:          nil,
				},
				InstanceType: "c5.xlarge",
				Count:        3,
				Tidb: &apiv1.Cluster_NodeSpec{
					ComponentType: common.ComponentType_TIDB,
					InstanceType:  "c5.xlarge",
					Count:         3,
				},
				Tiproxy: &apiv1.Cluster_NodeSpec{
					ComponentType: common.ComponentType_TIPROXY,
					InstanceType:  "c5.2xlarge",
					Count:         2,
				},
			},
		},
		{
			name: "with tiproxy in 'basic' group",
			fields: fields{
				KubeClient: fake.NewClientBuilder().WithObjects(
					newCluster("123", WithTiproxy("basic")),
				).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.GetTidbGroupRequest{
					TidbGroupId: "basic",
					ClusterId:   "456",
				},
			},
			wantErr: false,
			wantTg: &apiv1.TidbGroup{
				Id:                    "basic",
				ClusterId:             "456",
				IsDefault:             false,
				ClusterIdentifierHash: "clientidentifierhash2",
				EnablePublicLb:        true,
				EnablePrivateLb:       true,
				PublicLb: &apiv1.LoadBalancer{
					Host:            "public-extra-group",
					Port:            4000,
					CloudResourceId: "public-extra-group-id",
					IpAllowList:     sourceRangesToIpAllowList([]string{"192.168.0.1", "192.168.0.1"}),
					Synced:          ptr.To(true),
				},
				PrivateLb: &apiv1.LoadBalancer{
					Host:            "private-extra-group",
					Port:            4000,
					CloudResourceId: "private-extra-group-id",
					Synced:          ptr.To(true),
				},
				InstanceType: "c5.xlarge",
				Count:        3,
				Tidb: &apiv1.Cluster_NodeSpec{
					ComponentType: common.ComponentType_TIDB,
					InstanceType:  "c5.xlarge",
					Count:         3,
				},
				Tiproxy: &apiv1.Cluster_NodeSpec{
					ComponentType: common.ComponentType_TIPROXY,
					InstanceType:  "c5.2xlarge",
					Count:         2,
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &TidbGroupServer{
				KubeClient:    tt.fields.KubeClient,
				clusterReader: tt.fields.clusterReader,
			}
			gotTg, err := s.GetTidbGroup(tt.args.ctx, tt.args.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("TidbGroupServer.GetTidbGroup() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			assert.Equalf(t, tt.wantTg, gotTg, "TidbGroupServer.GetTidbGroup() = %v, want %v", gotTg, tt.wantTg)
		})
	}
}

func TestTidbGroupServer_ListTidbGroups(t *testing.T) {
	type fields struct {
		KubeClient    client.Client
		clusterReader internalutil.ClusterReader
	}
	type args struct {
		ctx context.Context
		req *apiv1.ListTidbGroupsRequest
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantRes *apiv1.ListTidbGroupsResponse
		wantErr bool
	}{
		{
			name: "basic",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123", withClusterReplica(0))).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.ListTidbGroupsRequest{
					ClusterId: "456",
				},
			},
			wantErr: false,
			wantRes: &apiv1.ListTidbGroupsResponse{
				TidbGroups: []*apiv1.TidbGroup{
					{
						Id:                    "123",
						ClusterId:             "456",
						IsDefault:             true,
						ClusterIdentifierHash: "clientidentifierhash",
						EnablePublicLb:        true,
						EnablePrivateLb:       true,
						PublicLb: &apiv1.LoadBalancer{
							Host:            "public",
							Port:            4000,
							CloudResourceId: "public-resource-id",
							IpAllowList:     sourceRangesToIpAllowList([]string{"0.0.0.0", "0.0.0.1"}),
							Synced:          nil,
						},
						PrivateLb: &apiv1.LoadBalancer{
							Host:            "private",
							Port:            4000,
							CloudResourceId: "private-resource-id",
							Synced:          nil,
						},
						InstanceType: "c5.xlarge",
						Count:        3,
						Tidb: &apiv1.Cluster_NodeSpec{
							ComponentType: common.ComponentType_TIDB,
							InstanceType:  "c5.xlarge",
							Count:         3,
						},
					},
					{
						Id:                    "basic",
						ClusterId:             "456",
						ClusterIdentifierHash: "clientidentifierhash2",
						EnablePublicLb:        true,
						EnablePrivateLb:       true,
						PublicLb: &apiv1.LoadBalancer{
							Host:            "public-extra-group",
							Port:            4000,
							CloudResourceId: "public-extra-group-id",
							IpAllowList:     sourceRangesToIpAllowList([]string{"192.168.0.1", "192.168.0.1"}),
							Synced:          ptr.To(true),
						},
						PrivateLb: &apiv1.LoadBalancer{
							Host:            "private-extra-group",
							Port:            4000,
							CloudResourceId: "private-extra-group-id",
							Synced:          ptr.To(true),
						},
						InstanceType: "c5.xlarge",
						Count:        3,
						Tidb: &apiv1.Cluster_NodeSpec{
							ComponentType: common.ComponentType_TIDB,
							InstanceType:  "c5.xlarge",
							Count:         3,
						},
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &TidbGroupServer{
				KubeClient:    tt.fields.KubeClient,
				clusterReader: tt.fields.clusterReader,
			}
			gotRes, err := s.ListTidbGroups(tt.args.ctx, tt.args.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("TidbGroupServer.ListTidbGroups() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			assert.Equalf(t, tt.wantRes, gotRes, "TidbGroupServer.ListTidbGroups() = %v, want %v", gotRes, tt.wantRes)
		})
	}
}

func TestTidbGroupServer_UpdateTidbGroup(t *testing.T) {
	type fields struct {
		KubeClient    client.Client
		clusterReader internalutil.ClusterReader
	}
	type args struct {
		ctx context.Context
		req *apiv1.UpdateTidbGroupRequest
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    *apiv1.TidbGroup
		wantErr bool
	}{
		{
			name: "basic",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.UpdateTidbGroupRequest{
					TidbGroupId: "basic",
					ClusterId:   "456",
					TidbGroup: &apiv1.TidbGroup{
						EnablePublicLb:  false,
						EnablePrivateLb: true,
					},
				},
			},
			wantErr: false,
			want: &apiv1.TidbGroup{
				EnablePrivateLb: true,
			},
		},
		{
			name: "default tidb group",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.UpdateTidbGroupRequest{
					TidbGroupId: "123",
					ClusterId:   "456",
					TidbGroup: &apiv1.TidbGroup{
						EnablePublicLb:  true,
						EnablePrivateLb: false,
					},
				},
			},
			wantErr: false,
			want: &apiv1.TidbGroup{
				EnablePublicLb: true,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &TidbGroupServer{
				KubeClient:    tt.fields.KubeClient,
				clusterReader: tt.fields.clusterReader,
			}
			got, err := s.UpdateTidbGroup(tt.args.ctx, tt.args.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("TidbGroupServer.UpdateTidbGroup() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("TidbGroupServer.UpdateTidbGroup() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestTidbGroupServer_SetPrivateLinkIPAddressType(t *testing.T) {
	tests := []struct {
		name       string
		groupID    string
		requested  apiv1.PrivateLinkIPAddressType
		nlb        apiv1.PrivateLinkIPAddressType
		setup      func(*cloudv1.Cluster)
		assertSpec func(*testing.T, *cloudv1.Cluster)
	}{
		{
			name:      "default group request and NLB become dualstack atomically",
			groupID:   "default-group",
			requested: apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_DUALSTACK,
			nlb:       apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_DUALSTACK,
			assertSpec: func(t *testing.T, cluster *cloudv1.Cluster) {
				assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeDualstack, cluster.Spec.SqlEndpoints.PrivateLinkIPAddressType)
				assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeDualstack, cluster.Spec.SqlEndpoints.PrivateLinkNLBIPAddressType)
				assert.Empty(t, cluster.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkIPAddressType)
			},
		},
		{
			name:      "extra group request changes without changing NLB target",
			groupID:   "basic",
			requested: apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_IPV4,
			setup: func(cluster *cloudv1.Cluster) {
				cluster.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkIPAddressType = cloudv1.PrivateLinkIPAddressTypeDualstack
				cluster.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkNLBIPAddressType = cloudv1.PrivateLinkIPAddressTypeDualstack
			},
			assertSpec: func(t *testing.T, cluster *cloudv1.Cluster) {
				assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeIPv4, cluster.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkIPAddressType)
				assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeDualstack, cluster.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkNLBIPAddressType)
				assert.Empty(t, cluster.Spec.SqlEndpoints.PrivateLinkIPAddressType)
			},
		},
		{
			name:    "NLB target changes without changing request target",
			groupID: "default-group",
			nlb:     apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_DUALSTACK,
			assertSpec: func(t *testing.T, cluster *cloudv1.Cluster) {
				assert.Empty(t, cluster.Spec.SqlEndpoints.PrivateLinkIPAddressType)
				assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeDualstack, cluster.Spec.SqlEndpoints.PrivateLinkNLBIPAddressType)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cluster := newCluster("default-group")
			cluster.Spec.Kubernetes.Provider = consts.ProviderTypeEKS
			if tt.setup != nil {
				tt.setup(cluster)
			}
			kubeClient := fake.NewClientBuilder().WithObjects(cluster).WithScheme(scheme.Scheme).Build()
			server := NewTidbGroupServer(kubeClient)
			req := &apiv1.SetPrivateLinkIPAddressTypeRequest{
				ClusterId:              cluster.Name,
				TidbGroupId:            tt.groupID,
				RequestedIpAddressType: tt.requested,
				NlbIpAddressType:       tt.nlb,
			}

			_, err := server.SetPrivateLinkIPAddressType(context.Background(), req)
			require.NoError(t, err)
			_, err = server.SetPrivateLinkIPAddressType(context.Background(), req)
			require.NoError(t, err)

			updated := &cloudv1.Cluster{}
			require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(cluster), updated))
			tt.assertSpec(t, updated)
		})
	}
}

func TestTidbGroupServer_SetPrivateLinkIPAddressTypeRetriesConflict(t *testing.T) {
	cluster := newCluster("default-group")
	cluster.Spec.Kubernetes.Provider = consts.ProviderTypeEKS
	kubeClient := fake.NewClientBuilder().WithObjects(cluster).WithScheme(scheme.Scheme).Build()
	conflictClient := &conflictOnceClient{Client: kubeClient}
	server := NewTidbGroupServer(conflictClient)

	_, err := server.SetPrivateLinkIPAddressType(context.Background(), &apiv1.SetPrivateLinkIPAddressTypeRequest{
		ClusterId:              cluster.Name,
		TidbGroupId:            "basic",
		RequestedIpAddressType: apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_DUALSTACK,
		NlbIpAddressType:       apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_DUALSTACK,
	})

	require.NoError(t, err)
	assert.Equal(t, 2, conflictClient.updateCalls)
	updated := &cloudv1.Cluster{}
	require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(cluster), updated))
	assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeDualstack, updated.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkIPAddressType)
	assert.Equal(t, cloudv1.PrivateLinkIPAddressTypeDualstack, updated.Spec.ExtraGroups[0].SqlEndpoints.PrivateLinkNLBIPAddressType)
}

func TestTidbGroupServer_SetPrivateLinkIPAddressTypeRejectsNonEKS(t *testing.T) {
	cluster := newCluster("default-group")
	cluster.Spec.Kubernetes.Provider = consts.ProviderTypeGKE
	kubeClient := fake.NewClientBuilder().WithObjects(cluster).WithScheme(scheme.Scheme).Build()
	server := NewTidbGroupServer(kubeClient)

	_, err := server.SetPrivateLinkIPAddressType(context.Background(), &apiv1.SetPrivateLinkIPAddressTypeRequest{
		ClusterId:              cluster.Name,
		TidbGroupId:            "default-group",
		RequestedIpAddressType: apiv1.PrivateLinkIPAddressType_PRIVATE_LINK_IP_ADDRESS_TYPE_DUALSTACK,
	})

	assert.Equal(t, codes.InvalidArgument, status.Code(err))
	latest := &cloudv1.Cluster{}
	require.NoError(t, kubeClient.Get(context.Background(), client.ObjectKeyFromObject(cluster), latest))
	assert.Empty(t, latest.Spec.SqlEndpoints.PrivateLinkIPAddressType)
}

type conflictOnceClient struct {
	client.Client
	updateCalls int
}

func (c *conflictOnceClient) Update(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
	c.updateCalls++
	if c.updateCalls == 1 {
		return apierrors.NewConflict(
			schema.GroupResource{Group: cloudv1.GroupVersion.Group, Resource: "clusters"},
			obj.GetName(),
			errors.New("injected conflict"),
		)
	}
	return c.Client.Update(ctx, obj, opts...)
}

func TestTidbGroupServer_SetPrivateLinkIPAddressTypeRejectsUnknownType(t *testing.T) {
	cluster := newCluster("default-group")
	kubeClient := fake.NewClientBuilder().WithObjects(cluster).WithScheme(scheme.Scheme).Build()
	server := NewTidbGroupServer(kubeClient)

	_, err := server.SetPrivateLinkIPAddressType(context.Background(), &apiv1.SetPrivateLinkIPAddressTypeRequest{
		ClusterId:              cluster.Name,
		TidbGroupId:            "default-group",
		RequestedIpAddressType: apiv1.PrivateLinkIPAddressType(3),
	})
	assert.Equal(t, codes.InvalidArgument, status.Code(err))

	latest := &cloudv1.Cluster{}
	require.NoError(t, server.KubeClient.Get(context.Background(), client.ObjectKeyFromObject(cluster), latest))
	assert.Empty(t, latest.Spec.SqlEndpoints.PrivateLinkIPAddressType)
}

func TestTidbGroupServer_SetPrivateLinkIPAddressTypeRejectsEmptyUpdate(t *testing.T) {
	cluster := newCluster("default-group")
	cluster.Spec.Kubernetes.Provider = consts.ProviderTypeEKS
	kubeClient := fake.NewClientBuilder().WithObjects(cluster).WithScheme(scheme.Scheme).Build()
	server := NewTidbGroupServer(kubeClient)

	_, err := server.SetPrivateLinkIPAddressType(context.Background(), &apiv1.SetPrivateLinkIPAddressTypeRequest{
		ClusterId:   cluster.Name,
		TidbGroupId: "default-group",
	})

	assert.Equal(t, codes.InvalidArgument, status.Code(err))
	latest := &cloudv1.Cluster{}
	require.NoError(t, server.KubeClient.Get(context.Background(), client.ObjectKeyFromObject(cluster), latest))
	assert.Empty(t, latest.Spec.SqlEndpoints.PrivateLinkIPAddressType)
	assert.Empty(t, latest.Spec.SqlEndpoints.PrivateLinkNLBIPAddressType)
}

func TestTidbGroupServer_DeleteTidbGroup(t *testing.T) {
	type fields struct {
		KubeClient    client.Client
		clusterReader internalutil.ClusterReader
	}
	type args struct {
		ctx context.Context
		req *apiv1.DeleteTidbGroupRequest
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    *apiv1.TidbGroup
		wantErr bool
	}{
		{
			name: "basic",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.DeleteTidbGroupRequest{
					TidbGroupId: "basic",
					ClusterId:   "456",
				},
			},
			wantErr: false,
			want: &apiv1.TidbGroup{
				Id:        "basic",
				ClusterId: "456",
			},
		},
		{
			name: "default tidb group",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx: nil,
				req: &apiv1.DeleteTidbGroupRequest{
					TidbGroupId: "123",
					ClusterId:   "456",
				},
			},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &TidbGroupServer{
				KubeClient:    tt.fields.KubeClient,
				clusterReader: tt.fields.clusterReader,
			}
			got, err := s.DeleteTidbGroup(tt.args.ctx, tt.args.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("TidbGroupServer.DeleteTidbGroup() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("TidbGroupServer.DeleteTidbGroup() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestTidbGroupServer_updateDNSIdentity(t *testing.T) {
	type fields struct {
		KubeClient    client.Client
		clusterReader internalutil.ClusterReader
	}
	type args struct {
		ctx                context.Context
		tidbGroupId        string
		requestedIdentity  string
		currentIdentity    string
		currentDnsDisabled bool
		checkSyncFunc      func() (*bool, *bool, error)
	}
	tests := []struct {
		name                    string
		fields                  fields
		args                    args
		wantErr                 bool
		wantErrCode             codes.Code
		wantIdentityUpdated     bool
		expectedDefaultIdentity string
		expectedExtraIdentity   string
	}{
		{
			name: "should skip update when identity already matches",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "123",
				requestedIdentity:  "same-identity",
				currentIdentity:    "same-identity",
				currentDnsDisabled: true,
				checkSyncFunc: func() (*bool, *bool, error) {
					return ptr.To(true), ptr.To(true), nil
				},
			},
			wantErr:             false,
			wantIdentityUpdated: false,
		},
		{
			name: "should fail when DNS is not disabled",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "123",
				requestedIdentity:  "new-identity",
				currentIdentity:    "old-identity",
				currentDnsDisabled: false, // DNS is enabled
				checkSyncFunc: func() (*bool, *bool, error) {
					return ptr.To(true), ptr.To(true), nil
				},
			},
			wantErr:             true,
			wantErrCode:         codes.FailedPrecondition,
			wantIdentityUpdated: false,
		},
		{
			name: "should fail when checkSyncFunc returns error",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "123",
				requestedIdentity:  "new-identity",
				currentIdentity:    "old-identity",
				currentDnsDisabled: true,
				checkSyncFunc: func() (*bool, *bool, error) {
					return nil, nil, status.Errorf(codes.Internal, "sync check failed")
				},
			},
			wantErr:             true,
			wantErrCode:         codes.Internal,
			wantIdentityUpdated: false,
		},
		{
			name: "should fail when sync status is nil",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "123",
				requestedIdentity:  "new-identity",
				currentIdentity:    "old-identity",
				currentDnsDisabled: true,
				checkSyncFunc: func() (*bool, *bool, error) {
					return nil, ptr.To(true), nil // public sync is nil
				},
			},
			wantErr:             true,
			wantErrCode:         codes.Internal,
			wantIdentityUpdated: false,
		},
		{
			name: "should fail when LoadBalancers not fully synced",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "123",
				requestedIdentity:  "new-identity",
				currentIdentity:    "old-identity",
				currentDnsDisabled: true,
				checkSyncFunc: func() (*bool, *bool, error) {
					return ptr.To(false), ptr.To(true), nil // public not synced
				},
			},
			wantErr:             true,
			wantErrCode:         codes.FailedPrecondition,
			wantIdentityUpdated: false,
		},
		{
			name: "should successfully update default group identity",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "123",
				requestedIdentity:  "new-default-identity",
				currentIdentity:    "old-identity",
				currentDnsDisabled: true,
				checkSyncFunc: func() (*bool, *bool, error) {
					return ptr.To(true), ptr.To(true), nil
				},
			},
			wantErr:                 false,
			wantIdentityUpdated:     true,
			expectedDefaultIdentity: "new-default-identity",
		},
		{
			name: "should successfully update extra group identity",
			fields: fields{
				KubeClient:    fake.NewClientBuilder().WithObjects(newCluster("123")).WithScheme(scheme.Scheme).Build(),
				clusterReader: &internalutil.LabelSelectorReader{},
			},
			args: args{
				ctx:                context.Background(),
				tidbGroupId:        "basic",
				requestedIdentity:  "new-extra-identity",
				currentIdentity:    "old-identity",
				currentDnsDisabled: true,
				checkSyncFunc: func() (*bool, *bool, error) {
					return ptr.To(true), ptr.To(true), nil
				},
			},
			wantErr:               false,
			wantIdentityUpdated:   true,
			expectedExtraIdentity: "new-extra-identity",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := &TidbGroupServer{
				KubeClient:    tt.fields.KubeClient,
				clusterReader: tt.fields.clusterReader,
			}

			// Get the cluster before update
			cluster, err := s.getCluster(tt.args.ctx, "456")
			assert.NoError(t, err)

			// Create a sugared logger for the test
			logger := tilog.WithContextS(tt.args.ctx)

			// Call the method under test
			err = s.updateDNSIdentity(
				tt.args.ctx,
				logger,
				cluster,
				tt.args.tidbGroupId,
				tt.args.requestedIdentity,
				tt.args.currentIdentity,
				tt.args.currentDnsDisabled,
				tt.args.checkSyncFunc,
			)

			// Check error expectations
			if tt.wantErr {
				assert.Error(t, err)
				if tt.wantErrCode != codes.OK {
					st, ok := status.FromError(err)
					assert.True(t, ok, "Expected gRPC status error")
					assert.Equal(t, tt.wantErrCode, st.Code())
				}
			} else {
				assert.NoError(t, err)
			}

			// Check identity update expectations
			if tt.wantIdentityUpdated {
				if isDefaultLoadBalancer(tt.args.tidbGroupId, cluster) {
					assert.Equal(t, tt.expectedDefaultIdentity, cluster.Spec.SqlEndpoints.DNS.Identity)
				} else {
					// Extra group
					for i, group := range cluster.Spec.ExtraGroups {
						if group.Name == tt.args.tidbGroupId {
							assert.Equal(t, tt.expectedExtraIdentity, cluster.Spec.ExtraGroups[i].SqlEndpoints.DNS.Identity)
						}
					}
				}
			}
		})
	}
}