package service

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"go.uber.org/mock/gomock"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/cloud/gcp"
)

const (
	testTcName    = "db"
	testNS        = "tidb12345"
	testClusterId = "12345"
)

func Test_validateRestartReq(t *testing.T) {
	tests := []struct {
		name    string
		req     *apiv1.RestartReq
		wantErr bool
	}{
		{
			name:    "nil request",
			req:     nil,
			wantErr: true,
		},
		{
			name: "empty target pod name",
			req: &apiv1.RestartReq{
				TargetType: apiv1.RestartTargetType_POD,
			},
			wantErr: true,
		},
		{
			name: "empty target node name",
			req: &apiv1.RestartReq{
				TargetType: apiv1.RestartTargetType_NODE,
			},
			wantErr: true,
		},
		{
			name: "empty target node name",
			req: &apiv1.RestartReq{
				TargetType: apiv1.RestartTargetType_COMPONENT,
			},
			wantErr: true,
		},
		{
			name: "invalid component name",
			req: &apiv1.RestartReq{
				TargetType: apiv1.RestartTargetType_COMPONENT,
				TargetName: "invalid_component",
			},
			wantErr: true,
		},
		{
			name: "valid pod",
			req: &apiv1.RestartReq{
				TargetType: apiv1.RestartTargetType_POD,
				TargetName: "valid_pod",
			},
			wantErr: false,
		},
		{
			name: "valid component",
			req: &apiv1.RestartReq{
				TargetType: apiv1.RestartTargetType_COMPONENT,
				TargetName: "pd",
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if err := validateRestartReq(tt.req); (err != nil) != tt.wantErr {
				t.Errorf("validateRestartReq() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func Test_restartPodSvc(t *testing.T) {
	tests := []struct {
		name          string
		podName       string
		useOperatorV2 bool
		force         bool
		pod           *v1.Pod
		wantErr       bool
	}{
		{
			name:    "pod does not exist",
			podName: "db-tidb-1",
			force:   false,
			wantErr: true,
		},
		{
			name:    "pod is deleting",
			podName: "db-tidb-1",
			force:   false,
			pod: &v1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:              "db-tidb-1",
					Namespace:         testNS,
					DeletionTimestamp: &metav1.Time{},
					Finalizers:        []string{"test"},
				},
			},
			wantErr: false,
		},
		{
			name:    "restart pd pod",
			podName: "db-pd-0",
			force:   false,
			pod: &v1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-pd-0",
					Namespace: testNS,
				},
			},
			wantErr: false,
		},
		{
			name:          "v2: restart pd pod",
			podName:       "db-pd-xxx",
			useOperatorV2: true,
			force:         false,
			pod: &v1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-pd-xxx",
					Namespace: testNS,
				},
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())

			var cli client.Client
			if tt.pod != nil {
				cli = fake.NewClientBuilder().WithScheme(s).WithObjects(tt.pod).Build()
			} else {
				cli = fake.NewClientBuilder().WithScheme(s).Build()
			}

			var opts []Option
			if tt.useOperatorV2 {
				opts = append(opts, WithTiDBOperatorVersion("v2"))
			} else {
				opts = append(opts, WithTiDBOperatorVersion("v1"))
			}
			svc := NewAgentService(cli, opts...)
			_, err := svc.restartPod(context.TODO(), &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   "12345",
				TargetType:  apiv1.RestartTargetType_POD,
				TargetName:  tt.podName,
				Force:       tt.force,
			})
			if (err != nil) != tt.wantErr {
				t.Errorf("restartPod() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func buildTestTC(name, namespace string, ready bool, apply func(*v1alpha1.TidbCluster)) *v1alpha1.TidbCluster {
	tc := &v1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        name,
			Namespace:   namespace,
			Annotations: map[string]string{},
			Labels: map[string]string{
				cloudv1.LabelUsedFor: cloudv1.UsedForCluster,
			},
		},
		Spec: v1alpha1.TidbClusterSpec{},
		Status: v1alpha1.TidbClusterStatus{
			Conditions: []v1alpha1.TidbClusterCondition{
				{
					Type:   v1alpha1.TidbClusterReady,
					Status: v1.ConditionFalse,
				},
			},
		},
	}

	if ready {
		tc.Status.Conditions[0].Status = v1.ConditionTrue
	}

	if apply != nil {
		apply(tc)
	}

	return tc
}

func buildComponentPod(name, ns, component string, apply func(pod *v1.Pod)) *v1.Pod {
	pod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels: map[string]string{
				componentLabelKey: component,
			},
		},
	}

	if apply != nil {
		apply(pod)
	}
	return pod
}

func Test_restartComponentSvc(t *testing.T) {
	tests := []struct {
		name          string
		componentName string
		force         bool
		tc            *v1alpha1.TidbCluster
		pods          []*v1.Pod
		wantErr       bool
		expect        func(*WithT, client.Client)
	}{
		{
			name:          "tc does not exist",
			componentName: "pd",
			force:         false,
			wantErr:       false,
		},
		{
			name:          "tc is not ready",
			componentName: "tikv",
			force:         false,
			tc: buildTestTC(tcName, testNS, false, func(tc *v1alpha1.TidbCluster) {
				tc.Spec.TiKV = &v1alpha1.TiKVSpec{}
			}),
			wantErr: true,
		},
		{
			name:          "graceful restart tikv",
			componentName: "tikv",
			force:         false,
			tc: buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
				tc.Spec.TiKV = &v1alpha1.TiKVSpec{}
			}),
			wantErr: false,
			expect: func(g *WithT, c client.Client) {
				var tc v1alpha1.TidbCluster
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: testNS,
					Name:      tcName,
				}, &tc)).To(Succeed())
				g.Expect(tc.Spec.TiKV.Annotations).To(HaveKey(operationIDAnnotationKey))
				g.Expect(tc.Spec.TiKV.Annotations[operationIDAnnotationKey]).To(Equal("test"))
			},
		},
		{
			name:          "restart tidb forcefully",
			componentName: "tidb",
			force:         true,
			tc: buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
				tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
			}),
			pods: []*v1.Pod{
				buildComponentPod("db-tidb-0", testNS, "tidb", nil),
				buildComponentPod("db-tidb-1", testNS, "tidb", nil),
				buildComponentPod("db-tikv-0", testNS, "tikv", nil),
				buildComponentPod("db-tidb-0", "tidb56789", "tidb", nil),
			},
			wantErr: false,
			expect: func(g *WithT, c client.Client) {
				var pod v1.Pod
				// should be deleted
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: testNS,
					Name:      "db-tidb-0",
				}, &pod)).ShouldNot(Succeed())
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: testNS,
					Name:      "db-tidb-1",
				}, &pod)).ShouldNot(Succeed())
				// should not be deleted
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: "tidb56789",
					Name:      "db-tidb-0",
				}, &pod)).Should(Succeed())
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: testNS,
					Name:      "db-tikv-0",
				}, &pod)).Should(Succeed())
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			g.Expect(v1alpha1.AddToScheme(s)).To(Succeed())

			cli := fake.NewClientBuilder().WithScheme(s).Build()
			if tt.tc != nil {
				g.Expect(cli.Create(context.TODO(), tt.tc)).To(Succeed())
			}
			for _, pod := range tt.pods {
				if pod != nil {
					g.Expect(cli.Create(context.TODO(), pod)).To(Succeed())
				}
			}

			svc := NewAgentService(cli)
			_, err := svc.restartComponent(context.TODO(), &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   testClusterId,
				TargetType:  apiv1.RestartTargetType_COMPONENT,
				TargetName:  tt.componentName,
				Force:       tt.force,
			})
			if (err != nil) != tt.wantErr {
				t.Errorf("restartComponent() error = %v, wantErr %v", err, tt.wantErr)
			}
			if tt.expect != nil {
				tt.expect(g, cli)
			}
		})
	}
}

func Test_restartTiDBClusterSvc(t *testing.T) {
	tests := []struct {
		name    string
		force   bool
		pods    []*v1.Pod
		tc      *v1alpha1.TidbCluster
		wantErr bool
		expect  func(*WithT, client.Client)
	}{
		{
			name:    "tc does not exist",
			force:   false,
			wantErr: true,
		},
		{
			name:    "tc is not ready",
			force:   false,
			tc:      buildTestTC(tcName, testNS, false, nil),
			wantErr: true,
		},
		{
			name:  "graceful restart tidb cluster",
			force: false,
			tc: buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
				tc.Spec.TiKV = &v1alpha1.TiKVSpec{}
				tc.Spec.PD = &v1alpha1.PDSpec{}
				tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
				tc.Spec.TiFlash = &v1alpha1.TiFlashSpec{}
				tc.Spec.TiProxy = &v1alpha1.TiProxySpec{}
			}),
			wantErr: false,
			expect: func(g *WithT, c client.Client) {
				var tc v1alpha1.TidbCluster
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: testNS,
					Name:      tcName,
				}, &tc)).To(Succeed())

				for _, comp := range tc.AllComponentSpec() {
					if comp.MemberType() != v1alpha1.DiscoveryMemberType {
						g.Expect(comp.Annotations()).To(HaveKey(operationIDAnnotationKey))
						g.Expect(comp.Annotations()[operationIDAnnotationKey]).To(Equal("test"))
					}
				}
			},
		},
		{
			name:  "force restart tidb cluster",
			force: true,
			tc: buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
				tc.Spec.TiKV = &v1alpha1.TiKVSpec{}
				tc.Spec.PD = &v1alpha1.PDSpec{}
				tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
			}),
			wantErr: false,
			pods: []*v1.Pod{
				buildComponentPod("db-tidb-0", testNS, "tidb", nil),
				buildComponentPod("db-pd-0", testNS, "pd", nil),
				buildComponentPod("db-tikv-0", testNS, "tikv", nil),
				buildComponentPod("db-tikv-0", "tidb56789", "tikv", nil),
				buildComponentPod("db-discovery", testNS, "discovery", nil),
			},
			expect: func(g *WithT, c client.Client) {
				var pod v1.Pod
				// should be deleted
				for _, name := range []string{"db-tidb-0", "db-pd-0", "db-tikv-0"} {
					g.Expect(c.Get(context.TODO(), client.ObjectKey{
						Namespace: testNS,
						Name:      name,
					}, &pod)).ShouldNot(Succeed())
				}
				// should not be deleted
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: testNS,
					Name:      "db-discovery",
				}, &pod)).Should(Succeed())
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Namespace: "tidb56789",
					Name:      "db-tikv-0",
				}, &pod)).Should(Succeed())
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			g.Expect(v1alpha1.AddToScheme(s)).To(Succeed())

			cli := fake.NewClientBuilder().WithScheme(s).Build()
			if tt.tc != nil {
				g.Expect(cli.Create(context.TODO(), tt.tc)).To(Succeed())
			}
			for _, pod := range tt.pods {
				if pod != nil {
					g.Expect(cli.Create(context.TODO(), pod)).To(Succeed())
				}
			}

			svc := NewAgentService(cli)
			_, err := svc.restartTiDBCluster(context.TODO(), &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   "12345",
				TargetType:  apiv1.RestartTargetType_CLUSTER,
				Force:       tt.force,
			})
			if (err != nil) != tt.wantErr {
				t.Errorf("restartTiDBCluster() error = %v, wantErr %v", err, tt.wantErr)
			}
			if tt.expect != nil {
				tt.expect(g, cli)
			}
		})
	}
}

func buildTestNode(name, providerID string, apply func(node *v1.Node)) *v1.Node {
	node := &v1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Spec: v1.NodeSpec{
			ProviderID: providerID,
		},
	}
	if apply != nil {
		apply(node)
	}
	return node
}

func Test_restartNodeSvc(t *testing.T) {
	tests := []struct {
		name       string
		provider   string
		targetName string
		force      bool
		pods       []*v1.Pod
		nodes      []*v1.Node
		wantErr    bool
		preExpect  func(*gomock.Controller, *WithT, *aws.MockClient, *gcp.MockClient)
		expect     func(*WithT, client.Client)
	}{
		{
			name:       "kind provider",
			provider:   "kind",
			targetName: "db-tidb-0",
			force:      false,
			wantErr:    true,
			preExpect: func(_ *gomock.Controller, g *WithT, awsCli *aws.MockClient, gcpCli *gcp.MockClient) {
				g.Expect(awsCli).To(BeNil())
				g.Expect(gcpCli).To(BeNil())
			},
		},
		{
			name:       "pod does not exist",
			provider:   "aws",
			targetName: "db-tidb-0",
			force:      false,
			wantErr:    true,
			preExpect: func(ctl *gomock.Controller, g *WithT, awsCli *aws.MockClient, gcpCli *gcp.MockClient) {
				g.Expect(gcpCli).To(BeNil())
				g.Expect(awsCli).NotTo(BeNil())
				awsCli.EXPECT().EC2(gomock.Any()).Times(0)
			},
		},
		{
			name:       "try to restart a non-tidb node",
			provider:   "aws",
			targetName: "db-discovery-foo",
			pods: []*v1.Pod{
				buildComponentPod("db-discovery-foo", testNS, "discovery", nil),
			},
			force:   false,
			wantErr: true,
			preExpect: func(_ *gomock.Controller, g *WithT, awsCli *aws.MockClient, gcpCli *gcp.MockClient) {
				g.Expect(gcpCli).To(BeNil())
				g.Expect(awsCli).NotTo(BeNil())
				awsCli.EXPECT().EC2(gomock.Any()).Times(0)
			},
		},
		{
			name:       "failed to get the node",
			provider:   "aws",
			targetName: "db-tidb-0",
			pods: []*v1.Pod{
				buildComponentPod("db-tidb-0", testNS, "tidb", func(pod *v1.Pod) {
					pod.Spec.NodeName = "node-0"
				}),
			},
			nodes:   []*v1.Node{},
			force:   false,
			wantErr: true,
			preExpect: func(_ *gomock.Controller, g *WithT, awsCli *aws.MockClient, gcpCli *gcp.MockClient) {
				g.Expect(gcpCli).To(BeNil())
				g.Expect(awsCli).NotTo(BeNil())
				awsCli.EXPECT().EC2(gomock.Any()).Times(0)
			},
		},
		{
			name:       "node has already been restarted",
			provider:   "aws",
			targetName: "db-tidb-0",
			pods: []*v1.Pod{
				buildComponentPod("db-tidb-0", testNS, "tidb", func(pod *v1.Pod) {
					pod.Spec.NodeName = "node-0"
				}),
			},
			nodes: []*v1.Node{
				buildTestNode("node-0", "", func(node *v1.Node) {
					node.Annotations = map[string]string{
						operationIDAnnotationKey: "test",
					}
				}),
			},
			force:   false,
			wantErr: false,
			preExpect: func(_ *gomock.Controller, g *WithT, awsCli *aws.MockClient, gcpCli *gcp.MockClient) {
				g.Expect(gcpCli).To(BeNil())
				g.Expect(awsCli).NotTo(BeNil())
				awsCli.EXPECT().EC2(gomock.Any()).Times(0)
			},
		},
		{
			name:       "restart GKE node",
			provider:   "gcp",
			targetName: "db-tidb-0",
			pods: []*v1.Pod{
				buildComponentPod("db-tidb-0", testNS, "tidb", func(pod *v1.Pod) {
					pod.Spec.NodeName = "node-0"
				}),
			},
			nodes: []*v1.Node{
				buildTestNode("node-0", "gce://test-project/us-west1-a/gke-xxx-node-po-000", nil),
				buildTestNode("node-1", "gce://test-project/us-west1-b/gke-xxx-node-po-001", nil),
			},
			force:   false,
			wantErr: false,
			preExpect: func(ctl *gomock.Controller, g *WithT, awsCli *aws.MockClient, gcpCli *gcp.MockClient) {
				g.Expect(gcpCli).NotTo(BeNil())
				g.Expect(awsCli).To(BeNil())
				gcpCli.EXPECT().RestartInstance(gomock.Any(), "test-project", "us-west1-a", "gke-xxx-node-po-000").Times(1)
			},
			expect: func(g *WithT, cli client.Client) {
				var node v1.Node
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{
					Name: "node-0",
				}, &node)).To(Succeed())
				g.Expect(node.Annotations).To(HaveKeyWithValue(operationIDAnnotationKey, "test"))

				g.Expect(cli.Get(context.TODO(), client.ObjectKey{
					Name: "node-1",
				}, &node)).To(Succeed())
				g.Expect(node.Annotations).NotTo(HaveKeyWithValue(operationIDAnnotationKey, "test"))
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())

			objs := make([]client.Object, 0, len(tt.pods)+len(tt.nodes))
			for _, pod := range tt.pods {
				if pod != nil {
					objs = append(objs, pod)
				}
			}
			for _, node := range tt.nodes {
				if node != nil {
					objs = append(objs, node)
				}
			}
			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()

			ctl := gomock.NewController(t)
			defer ctl.Finish()

			var (
				svc        *AgentService
				mockAWSCli *aws.MockClient
				mockGCPCli *gcp.MockClient
			)
			switch tt.provider {
			case "aws":
				mockAWSCli = aws.NewMockClient(ctl)
				nr := NewAWSNodeRestarter(mockAWSCli)
				svc = NewAgentService(cli, WithNodeRestarter(nr))
			case "gcp":
				mockGCPCli = gcp.NewMockClient(ctl)
				nr := NewGCPNodeRestarter(mockGCPCli)
				svc = NewAgentService(cli, WithNodeRestarter(nr))
			case "kind":
				nr := NewKindNodeRestarter()
				svc = NewAgentService(cli, WithNodeRestarter(nr))
			case "azure":
				// TODO: add azure control logic
				svc = NewAgentService(cli)
			default:
				t.Errorf("unknown provider: %s", tt.provider)
			}

			if tt.preExpect != nil {
				tt.preExpect(ctl, g, mockAWSCli, mockGCPCli)
			}

			_, err := svc.restartNode(context.TODO(), &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   "12345",
				TargetType:  apiv1.RestartTargetType_NODE,
				TargetName:  tt.targetName,
				Force:       tt.force,
			})
			if (err != nil) != tt.wantErr {
				t.Errorf("restartNode() error = %v, wantErr %v", err, tt.wantErr)
			}
			if tt.expect != nil {
				tt.expect(g, cli)
			}
		})
	}
}

func buildTestStatus(ready, synced bool) opv2.CommonStatus {
	return opv2.CommonStatus{
		Conditions: []metav1.Condition{
			{
				Type:   opv2.CondReady,
				Status: boolToConditionStatus(ready),
			},
			{
				Type:   opv2.CondSynced,
				Status: boolToConditionStatus(synced),
			},
		},
	}
}

func boolToConditionStatus(status bool) metav1.ConditionStatus {
	if status {
		return metav1.ConditionTrue
	}
	return metav1.ConditionFalse
}

func TestAgentService_graceRestartAllGroupsOfComponent4V1(t *testing.T) {
	tests := map[string]struct {
		req     *apiv1.RestartReq
		tcs     []*v1alpha1.TidbCluster
		wantErr bool
		expect  func(*WithT, client.Client)
	}{
		"restart pd component in single TC": {
			req: &apiv1.RestartReq{
				OperationId: "test-op",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.PD = &v1alpha1.PDSpec{}
				}),
			},
			wantErr: false,
			expect: func(g *WithT, cli client.Client) {
				var tc v1alpha1.TidbCluster
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: tcName}, &tc)).To(Succeed())
				g.Expect(tc.Spec.PD.Annotations).To(HaveKeyWithValue(operationIDAnnotationKey, "test-op"))
			},
		},
		"restart tidb component with multiple TCs": {
			req: &apiv1.RestartReq{
				OperationId: "multi-tc-op",
				ClusterId:   testClusterId,
				TargetName:  "tidb",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
				}),
				buildTestTC("db-group1", testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
				}),
			},
			wantErr: false,
			expect: func(g *WithT, cli client.Client) {
				var tc v1alpha1.TidbCluster
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: tcName}, &tc)).To(Succeed())
				g.Expect(tc.Spec.TiDB.Annotations).To(HaveKeyWithValue(operationIDAnnotationKey, "multi-tc-op"))

				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: "db-group1"}, &tc)).To(Succeed())
				g.Expect(tc.Spec.TiDB.Annotations).To(HaveKeyWithValue(operationIDAnnotationKey, "multi-tc-op"))
			},
		},
		"restart specific target group": {
			req: &apiv1.RestartReq{
				OperationId: "target-group-op",
				ClusterId:   testClusterId,
				TargetName:  "tidb",
				TargetGroup: "db-group1",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
				}),
				buildTestTC("db-group1", testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiDB = &v1alpha1.TiDBSpec{}
				}),
			},
			wantErr: false,
			expect: func(g *WithT, cli client.Client) {
				var tc v1alpha1.TidbCluster
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: tcName}, &tc)).To(Succeed())
				g.Expect(tc.Spec.TiDB.Annotations).NotTo(HaveKey(operationIDAnnotationKey))

				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: "db-group1"}, &tc)).To(Succeed())
				g.Expect(tc.Spec.TiDB.Annotations).To(HaveKeyWithValue(operationIDAnnotationKey, "target-group-op"))
			},
		},
		"target group not found": {
			req: &apiv1.RestartReq{
				OperationId: "missing-group-op",
				ClusterId:   testClusterId,
				TargetName:  "tiflash",
				TargetGroup: "nonexistent-group",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiFlash = &v1alpha1.TiFlashSpec{}
				}),
			},
			wantErr: true,
		},
		"empty TC list": {
			req: &apiv1.RestartReq{
				OperationId: "empty-list-op",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			tcs:     []*v1alpha1.TidbCluster{},
			wantErr: false,
		},
		"TC not ready should fail": {
			req: &apiv1.RestartReq{
				OperationId: "not-ready-op",
				ClusterId:   testClusterId,
				TargetName:  "tiproxy",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, false, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiProxy = &v1alpha1.TiProxySpec{}
				}),
			},
			wantErr: true,
		},
		"no update no error when tc is NOT ready but operationId already set": {
			req: &apiv1.RestartReq{
				OperationId: "existing-op-id",
				ClusterId:   testClusterId,
				TargetName:  "tikv",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, false, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.TiKV = &v1alpha1.TiKVSpec{
						ComponentSpec: v1alpha1.ComponentSpec{
							Annotations: map[string]string{
								operationIDAnnotationKey: "existing-op-id",
							},
						},
					}
				}),
			},
			wantErr: false,
			expect: func(g *WithT, cli client.Client) {
				var tc v1alpha1.TidbCluster
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: tcName}, &tc)).To(Succeed())
				g.Expect(tc.Spec.TiKV.Annotations[operationIDAnnotationKey]).To(Equal("existing-op-id"))
			},
		},
		"mixed ready and not ready TCs": {
			req: &apiv1.RestartReq{
				OperationId: "mixed-op",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			tcs: []*v1alpha1.TidbCluster{
				buildTestTC(tcName, testNS, true, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.PD = &v1alpha1.PDSpec{}
				}),
				buildTestTC("db-group1", testNS, false, func(tc *v1alpha1.TidbCluster) {
					tc.Spec.PD = &v1alpha1.PDSpec{}
				}),
			},
			wantErr: true,
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			g.Expect(v1alpha1.AddToScheme(s)).To(Succeed())

			objs := make([]client.Object, 0, len(tt.tcs))
			for _, tc := range tt.tcs {
				objs = append(objs, tc)
			}
			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()

			svc := NewAgentService(cli, WithTiDBOperatorVersion("v1"))
			_, err := svc.graceRestartAllGroupsOfComponent4V1(context.TODO(), tt.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("restartAllGroupsOfComponent4V1() error = %v, wantErr %v", err, tt.wantErr)
			}
			if tt.expect != nil {
				tt.expect(g, cli)
			}
		})
	}
}

func TestAgentService_graceRestartAllGroupsOfComponent4V2(t *testing.T) {
	tests := []struct {
		name    string
		req     *apiv1.RestartReq
		group   client.Object
		wantErr bool
		expect  func(*WithT, client.Client)
	}{
		{
			name: "restart pd group when it's ready and synced",
			req: &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			group: &opv2.PDGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db",
					Namespace: testNS,
				},
				Status: opv2.PDGroupStatus{
					CommonStatus: buildTestStatus(true, true),
				},
			},
			expect: func(g *WithT, cli client.Client) {
				var pdg opv2.PDGroup
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: "db"}, &pdg)).To(Succeed())
				g.Expect(pdg.Spec.Template.Annotations).To(HaveKeyWithValue(operationIDAnnotationKey, "test"))
			},
		},
		{
			name: "can not restart pd group when it's not ready",
			req: &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			group: &opv2.PDGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db",
					Namespace: testNS,
				},
				Status: opv2.PDGroupStatus{
					CommonStatus: buildTestStatus(false, true),
				},
			},
			wantErr: true,
		},
		{
			name: "can not restart pd group when it's not synced",
			req: &apiv1.RestartReq{
				OperationId: "test",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			group: &opv2.PDGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db",
					Namespace: testNS,
				},
				Status: opv2.PDGroupStatus{
					CommonStatus: buildTestStatus(true, false),
				},
			},
			wantErr: true,
		},
		{
			name: "tikv group not ready should fail",
			req: &apiv1.RestartReq{
				OperationId: "tikv-test",
				ClusterId:   testClusterId,
				TargetName:  "tikv",
			},
			group: &opv2.TiKVGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db",
					Namespace: testNS,
				},
				Status: opv2.TiKVGroupStatus{
					CommonStatus: buildTestStatus(false, true),
				},
			},
			wantErr: true,
		},
		{
			name: "unsupported component type",
			req: &apiv1.RestartReq{
				OperationId: "invalid-test",
				ClusterId:   testClusterId,
				TargetName:  "unknown",
			},
			group:   &opv2.TiKVGroup{},
			wantErr: true,
		},
		{
			name: "no update when operationId matches",
			req: &apiv1.RestartReq{
				OperationId: "existing-op",
				ClusterId:   testClusterId,
				TargetName:  "pd",
			},
			group: &opv2.PDGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db",
					Namespace: testNS,
				},
				Spec: opv2.PDGroupSpec{
					Template: opv2.PDTemplate{
						ObjectMeta: opv2.ObjectMeta{
							Annotations: map[string]string{operationIDAnnotationKey: "existing-op"},
						},
					},
				},
				Status: opv2.PDGroupStatus{
					CommonStatus: buildTestStatus(true, true),
				},
			},
			expect: func(g *WithT, cli client.Client) {
				var pdg opv2.PDGroup
				g.Expect(cli.Get(context.TODO(), client.ObjectKey{Namespace: testNS, Name: "db"}, &pdg)).To(Succeed())
				g.Expect(pdg.Spec.Template.Annotations[operationIDAnnotationKey]).To(Equal("existing-op"))
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(opv2.AddToScheme(s)).To(Succeed())

			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(tt.group).Build()
			svc := NewAgentService(cli)
			_, err := svc.graceRestartAllGroupsOfComponent4V2(context.TODO(), tt.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("restartComponentGroup() error = %v, wantErr %v", err, tt.wantErr)
			}
			if tt.expect != nil {
				tt.expect(g, cli)
			}
		})
	}
}