package operation

import (
	"context"
	"testing"
	"time"

	"github.com/crossplane-contrib/provider-aws/apis/eks/v1beta1"
	. "github.com/onsi/gomega"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"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/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gkemindomains"
	utiltime "github.com/tidbcloud/infra-provider/pkg/util/time"
)

const (
	testClusterId = "12345"
	testNS        = "test"
	testEKSName   = "test-eks"
)

func TestServer_UpgradeNodes(t *testing.T) {
	tests := []struct {
		name            string
		existingUO      *cloudv1.UpgradeOperation
		existingCluster *cloudv1.Cluster
		req             *apiv1.UpgradeNodesRequest
		now             *time.Time
		wantResp        *apiv1.UpgradeNodesResponse
		wantErr         bool
		expectUO        *cloudv1.UpgradeOperation
		checkEndTime    bool // If true, check that EndTime is set
	}{
		{
			name: "cluster not found - should return EndTime and no error",
			req: &apiv1.UpgradeNodesRequest{
				IdempotencyKey: "test-ns/test-uo",
				ClusterId:      testClusterId,
				NodeType:       apiv1.UpgradeNodesRequest_TIDB,
				DesiredVersion: "1.29",
				Paused:         false,
				Options: map[string]string{
					"scale-out-extra-tidb": "true",
				},
			},
			now:          ptr.To(time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)),
			wantErr:      false, // Changed: should not return error when cluster not found
			checkEndTime: true,  // Verify EndTime is set when cluster not found
		},
		{
			name:            "create CR to upgrade tidb nodes",
			existingCluster: util.BuildTestClusterCR(testClusterId, testNS, nil),
			req: &apiv1.UpgradeNodesRequest{
				IdempotencyKey: "upgrade-tidb",
				ClusterId:      testClusterId,
				NodeType:       apiv1.UpgradeNodesRequest_TIDB,
				DesiredVersion: "1.29",
				Paused:         false,
				Options: map[string]string{
					"scale-out-extra-tidb":              "true",
					gkemindomains.UpgradeDecisionOption: "min-domains",
				},
			},
			now:      ptr.To(time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)),
			wantResp: &apiv1.UpgradeNodesResponse{},
			expectUO: &cloudv1.UpgradeOperation{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "upgrade-tidb-12345",
					Namespace: testNS,
				},
				Spec: cloudv1.UpgradeOperationSpec{
					Schedule: time.Date(2024, 1, 1, 0, 0, 30, 0, time.UTC).Format(cloudv1.UpgradeOperationScheduleLayout),
					Paused:   false,
					Steps: []cloudv1.Step{
						{
							Target:  cloudv1.TiDBNodes,
							Version: "1.29",
							Cluster: testClusterId,
							Options: map[string]string{
								"scale-out-extra-tidb": "true",
							},
						},
					},
				},
			},
		},
		{
			name:            "update CR to pause upgrading tidb nodes",
			existingCluster: util.BuildTestClusterCR(testClusterId, testNS, nil),
			existingUO: &cloudv1.UpgradeOperation{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "upgrade-tidb-12345",
					Namespace: testNS,
				},
				Spec: cloudv1.UpgradeOperationSpec{
					Paused: false,
				},
			},
			req: &apiv1.UpgradeNodesRequest{
				IdempotencyKey: "upgrade-tidb",
				ClusterId:      testClusterId,
				NodeType:       apiv1.UpgradeNodesRequest_TIDB,
				DesiredVersion: "1.29",
				Paused:         true,
			},
			wantResp: &apiv1.UpgradeNodesResponse{},
			expectUO: &cloudv1.UpgradeOperation{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "upgrade-tidb-12345",
					Namespace: testNS,
				},
				Spec: cloudv1.UpgradeOperationSpec{
					Paused: true,
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			s := runtime.NewScheme()
			g.Expect(cloudv1.AddToScheme(s)).Should(Succeed())

			var objs []client.Object
			if tt.existingUO != nil {
				objs = append(objs, tt.existingUO)
			}
			if tt.existingCluster != nil {
				objs = append(objs, tt.existingCluster)
			}
			kubeClient := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()

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

			mockClock := utiltime.NewMockClock(ctl)
			if tt.now != nil {
				mockClock.EXPECT().Now().Return(*tt.now).Times(1)
			}

			got, err := NewServer(kubeClient, mockClock).UpgradeNodes(context.TODO(), tt.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("UpgradeNodes() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			// Check EndTime is set when cluster not found
			if tt.checkEndTime {
				g.Expect(got).ShouldNot(BeNil(), "Response should not be nil when cluster not found")
				g.Expect(got.EndTime).ShouldNot(BeNil(), "EndTime should be set when cluster not found")
				// Verify it's a reasonable time (around the mocked Now time)
				if tt.now != nil {
					g.Expect(got.EndTime.AsTime()).Should(BeTemporally("~", *tt.now, time.Second))
				}
			} else if tt.wantResp != nil {
				g.Expect(got).Should(Equal(tt.wantResp))
			}

			if tt.expectUO != nil {
				var uo cloudv1.UpgradeOperation
				err = kubeClient.Get(context.TODO(), client.ObjectKey{Name: tt.expectUO.Name, Namespace: tt.expectUO.Namespace}, &uo)
				g.Expect(err).Should(BeNil())
				g.Expect(uo.Spec.Schedule).Should(Equal(tt.expectUO.Spec.Schedule))
				g.Expect(uo.Spec.Steps).Should(Equal(tt.expectUO.Spec.Steps))
				g.Expect(uo.Spec.Paused).Should(Equal(tt.expectUO.Spec.Paused))
			}
		})
	}
}

func TestMakeUpgradeStepSanitizesOptionsWithoutMutatingRequest(t *testing.T) {
	g := NewWithT(t)
	req := &apiv1.UpgradeNodesRequest{
		ClusterId:      testClusterId,
		NodeType:       apiv1.UpgradeNodesRequest_TIDB,
		DesiredVersion: "1.30",
		Options: map[string]string{
			"scale-out-extra-tidb":              "true",
			gkemindomains.UpgradeDecisionOption: "legacy",
		},
	}

	step := makeUpgradeStep(req)

	g.Expect(step.Options).Should(Equal(map[string]string{"scale-out-extra-tidb": "true"}))
	g.Expect(req.Options).Should(Equal(map[string]string{
		"scale-out-extra-tidb":              "true",
		gkemindomains.UpgradeDecisionOption: "legacy",
	}))
}

func TestServer_CreateOperation(t *testing.T) {
	tests := []struct {
		name            string
		existingCRs     []client.Object
		existingCluster *cloudv1.Cluster
		req             *apiv1.CreateOperationRequest
		wantResp        *apiv1.CreateOperationResponse
		wantErr         bool
	}{
		{
			name: "Shutdown upgrade nodes: k8s control plane is not upgraded yet",
			existingCluster: util.BuildTestClusterCR(testClusterId, testNS, func(cluster *cloudv1.Cluster) {
				cluster.Spec.Kubernetes.Provider = cloudv1.ProviderTypeEKS
				cluster.Spec.Kubernetes.Name = testEKSName
			}),
			existingCRs: []client.Object{
				&v1beta1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: testEKSName},
					Status:     v1beta1.ClusterStatus{AtProvider: v1beta1.ClusterObservation{Version: "1.29"}},
				},
			},
			req: &apiv1.CreateOperationRequest{
				ClusterId: testClusterId,
				Operation: &apiv1.Operation{
					IdempotencyKey: "test",
					Type:           apiv1.Operation_SHUTDOWN_UPGRADE_NODES,
					ClusterId:      testClusterId,
					Args: &apiv1.Operation_ShutdownUpgradeNodes{
						ShutdownUpgradeNodes: &apiv1.Operation_ShutDownUpgradeNodes{
							DesiredVersion: "1.30",
						},
					},
				},
			},
			wantErr: true,
		},
		{
			name: "Shutdown upgrade nodes: happy path",
			existingCluster: util.BuildTestClusterCR(testClusterId, testNS, func(cluster *cloudv1.Cluster) {
				cluster.Spec.Kubernetes.Provider = cloudv1.ProviderTypeEKS
				cluster.Spec.Kubernetes.Name = testEKSName
			}),
			existingCRs: []client.Object{
				&v1beta1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: testEKSName},
					Status:     v1beta1.ClusterStatus{AtProvider: v1beta1.ClusterObservation{Version: "1.30"}},
				},
			},
			req: &apiv1.CreateOperationRequest{
				ClusterId: testClusterId,
				Operation: &apiv1.Operation{
					IdempotencyKey: "test",
					Type:           apiv1.Operation_SHUTDOWN_UPGRADE_NODES,
					ClusterId:      testClusterId,
					Args: &apiv1.Operation_ShutdownUpgradeNodes{
						ShutdownUpgradeNodes: &apiv1.Operation_ShutDownUpgradeNodes{
							DesiredVersion: "1.30",
						},
					},
				},
			},
			wantResp: &apiv1.CreateOperationResponse{Id: "shutdown-upgrade-nodes-test"},
			wantErr:  false,
		},
		{
			name:            "Offline AZ: happy path",
			existingCluster: util.BuildTestClusterCR(testClusterId, testNS, nil),
			req: &apiv1.CreateOperationRequest{
				ClusterId: testClusterId,
				Operation: &apiv1.Operation{
					IdempotencyKey: "offline-az",
					Type:           apiv1.Operation_OFFLINE_AZ,
					ClusterId:      testClusterId,
					Args: &apiv1.Operation_OfflineAz{
						OfflineAz: &apiv1.Operation_OfflineAZ{
							AvailabilityZone: "us-east-1a",
						},
					},
				},
			},
			wantResp: &apiv1.CreateOperationResponse{Id: "offline-az-offline-az"},
			wantErr:  false,
		},
		{
			name:            "Recover AZ: happy path",
			existingCluster: util.BuildTestClusterCR(testClusterId, testNS, nil),
			req: &apiv1.CreateOperationRequest{
				ClusterId: testClusterId,
				Operation: &apiv1.Operation{
					IdempotencyKey: "recover-az",
					Type:           apiv1.Operation_RECOVER_AZ,
					ClusterId:      testClusterId,
					Args: &apiv1.Operation_RecoverAz{
						RecoverAz: &apiv1.Operation_RecoverAZ{
							AvailabilityZone: "us-east-1a",
						},
					},
				},
			},
			wantResp: &apiv1.CreateOperationResponse{Id: "recover-az-recover-az"},
			wantErr:  false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			s := runtime.NewScheme()
			g.Expect(cloudv1.AddToScheme(s)).Should(Succeed())
			g.Expect(v1beta1.SchemeBuilder.AddToScheme(s)).Should(Succeed())

			objs := make([]client.Object, 0, len(tt.existingCRs)+1)
			objs = append(objs, tt.existingCRs...)
			if tt.existingCluster != nil {
				objs = append(objs, tt.existingCluster)
			}
			kubeClient := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()

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

			got, err := NewServer(kubeClient, utiltime.RealClock{}).CreateOperation(context.TODO(), tt.req)
			if (err != nil) != tt.wantErr {
				t.Errorf("CreateOperation() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			g.Expect(got).Should(Equal(tt.wantResp))

			if got != nil {
				var wf cloudv1.Workflow
				err = kubeClient.Get(context.TODO(), client.ObjectKey{Name: got.Id, Namespace: testNS}, &wf)
				g.Expect(err).Should(BeNil())
				switch tt.req.Operation.Type {
				case apiv1.Operation_OFFLINE_AZ:
					g.Expect(wf.Spec.Steps).To(HaveLen(1))
					g.Expect(wf.Spec.Steps[0].Uses).To(Equal(consts.WfActionOfflineAZAction))
					g.Expect(wf.Spec.Steps[0].Inputs.Object["availabilityZone"]).To(Equal("us-east-1a"))
				case apiv1.Operation_RECOVER_AZ:
					g.Expect(wf.Spec.Steps).To(HaveLen(1))
					g.Expect(wf.Spec.Steps[0].Uses).To(Equal(consts.WfActionRecoverAZAction))
					g.Expect(wf.Spec.Steps[0].Inputs.Object["availabilityZone"]).To(Equal("us-east-1a"))
				}
			}
		})
	}
}

func TestBuildMigrateTiKVToDualDiskWithActiveGKEVAC(t *testing.T) {
	tests := []struct {
		name        string
		provider    string
		storageType string
		currentVAC  string
		modifiedVAC string
		enableVAC   bool
		clearStatus bool
		wantCode    codes.Code
	}{
		{
			name:        "reject non-Hyperdisk volume with current VAC",
			provider:    cloudv1.ProviderTypeGKE,
			storageType: consts.GCPVolumeTypePDSSD,
			currentVAC:  "tikv-hyperdisk-balanced-3000-140",
			wantCode:    codes.FailedPrecondition,
		},
		{
			name:        "reject non-Hyperdisk volume with modifying VAC",
			provider:    cloudv1.ProviderTypeGKE,
			storageType: consts.GCPVolumeTypePDSSD,
			modifiedVAC: "tikv-hyperdisk-balanced-6000-140",
			wantCode:    codes.FailedPrecondition,
		},
		{
			name:        "wait for volume status when VAC is enabled",
			provider:    cloudv1.ProviderTypeGKE,
			storageType: consts.GCPVolumeTypePDSSD,
			enableVAC:   true,
			clearStatus: true,
			wantCode:    codes.Unavailable,
		},
		{
			name:        "allow missing volume status when VAC is disabled",
			provider:    cloudv1.ProviderTypeGKE,
			storageType: consts.GCPVolumeTypePDSSD,
			clearStatus: true,
			wantCode:    codes.OK,
		},
		{
			name:        "allow Hyperdisk volume with active VAC",
			provider:    cloudv1.ProviderTypeGKE,
			storageType: consts.GCPVolumeTypeHyperdiskBalanced,
			currentVAC:  "tikv-hyperdisk-balanced-3000-140",
			wantCode:    codes.OK,
		},
		{
			name:        "allow non-Hyperdisk volume without active VAC",
			provider:    cloudv1.ProviderTypeGKE,
			storageType: consts.GCPVolumeTypePDSSD,
			wantCode:    codes.OK,
		},
		{
			name:        "do not affect EKS",
			provider:    cloudv1.ProviderTypeEKS,
			storageType: consts.GCPVolumeTypePDSSD,
			currentVAC:  "tikv-gp3-3000-125",
			wantCode:    codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cluster := util.BuildTestClusterCR(testClusterId, testNS, func(cluster *cloudv1.Cluster) {
				cluster.Labels[cloudv1.ProviderTypeLabel] = tt.provider
				if tt.enableVAC {
					cluster.EnableFeature(consts.ClusterFeatureUseVolumeAttributesClass)
				}
				cluster.Status.DB.TiKV = &cloudv1.TiKVStatus{Volume: &cloudv1.VolumeStatus{
					CurrentVolumeAttributesClass:  tt.currentVAC,
					ModifiedVolumeAttributesClass: tt.modifiedVAC,
				}}
				if tt.clearStatus {
					cluster.Status.DB.TiKV.Volume = nil
				}
			})
			args := &apiv1.Operation_MigrateTiKVToDualDisk{
				AdditionalStorage: &apiv1.Operation_MigrateTiKVToDualDisk_AdditionalStorage{Type: tt.storageType},
			}
			wf := &cloudv1.Workflow{}

			err := (&Server{}).buildMigrateTiKVToDualDiskWF(wf, args, cluster)
			if got := status.Code(err); got != tt.wantCode {
				t.Fatalf("buildMigrateTiKVToDualDiskWF() code = %s, want %s, err = %v", got, tt.wantCode, err)
			}
			if tt.wantCode != codes.OK && len(wf.Spec.Steps) != 0 {
				t.Fatalf("invalid migration created %d workflow steps", len(wf.Spec.Steps))
			}
		})
	}
}

func TestServer_canUpgradeNodes(t *testing.T) {
	tests := []struct {
		name        string
		existingCRs []client.Object
		cluster     *cloudv1.Cluster
		args        *apiv1.Operation_ShutDownUpgradeNodes
		want        bool
		wantErr     bool
	}{
		{
			name: "happy path",
			existingCRs: []client.Object{
				&v1beta1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: testEKSName},
					Status:     v1beta1.ClusterStatus{AtProvider: v1beta1.ClusterObservation{Version: "1.30"}},
				},
			},
			args: &apiv1.Operation_ShutDownUpgradeNodes{
				DesiredVersion: "1.30",
			},
			want: true,
		},
		{
			name: "controlplane's version is newer than nodes' version",
			existingCRs: []client.Object{
				&v1beta1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: testEKSName},
					Status:     v1beta1.ClusterStatus{AtProvider: v1beta1.ClusterObservation{Version: "1.30"}},
				},
			},
			args: &apiv1.Operation_ShutDownUpgradeNodes{
				DesiredVersion: "1.29",
			},
			want: true,
		},
		{
			name: "controlplane's version is older than nodes' version",
			existingCRs: []client.Object{
				&v1beta1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: testEKSName},
					Status:     v1beta1.ClusterStatus{AtProvider: v1beta1.ClusterObservation{Version: "1.28"}},
				},
			},
			args: &apiv1.Operation_ShutDownUpgradeNodes{
				DesiredVersion: "1.30",
			},
			want: false,
		},
		{
			name:        "no eks CR",
			existingCRs: []client.Object{},
			args: &apiv1.Operation_ShutDownUpgradeNodes{
				DesiredVersion: "1.30",
			},
			want:    false,
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			s := runtime.NewScheme()
			g.Expect(cloudv1.AddToScheme(s)).Should(Succeed())
			g.Expect(v1beta1.SchemeBuilder.AddToScheme(s)).Should(Succeed())

			objs := make([]client.Object, 0, len(tt.existingCRs))
			objs = append(objs, tt.existingCRs...)
			kubeClient := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()

			cluster := util.BuildTestClusterCR(testClusterId, testNS, func(cluster *cloudv1.Cluster) {
				cluster.Spec.Kubernetes.Provider = cloudv1.ProviderTypeEKS
				cluster.Spec.Kubernetes.Name = testEKSName
			})
			server := NewServer(kubeClient, utiltime.RealClock{})
			got, err := server.canUpgradeNodes(context.TODO(), cluster, tt.args)
			if (err != nil) != tt.wantErr {
				t.Errorf("canUpgradeNodes() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("canUpgradeNodes() got = %v, want %v", got, tt.want)
			}
		})
	}
}