package gke

import (
	"context"
	"fmt"
	"reflect"
	"testing"

	"cloud.google.com/go/container/apiv1/containerpb"
	gkev1beta1 "github.com/crossplane-contrib/provider-gcp/apis/container/v1beta1"
	"github.com/go-logr/logr"
	"github.com/google/go-cmp/cmp"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/pointer"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/nodepool"
)

func Test_filterLegacyNodePoolsByAntiLabels(t *testing.T) {
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "c1",
			Namespace: "ns1",
		},
	}
	currentVersion := "1.28"
	desiredVersion := "1.29"
	machineType := "n2-standard-4"
	currentName := nodepool.GetGKENodePoolName(cluster, machineType, currentVersion)
	desiredName := nodepool.GetGKENodePoolName(cluster, machineType, desiredVersion)

	makeNP := func(name string, metaLabels, specLabels map[string]string, mt *string) gkev1beta1.NodePool {
		np := gkev1beta1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:   name,
				Labels: metaLabels,
			},
		}
		if mt != nil || specLabels != nil {
			np.Spec.ForProvider.Config = &gkev1beta1.NodeConfig{
				MachineType: mt,
				Labels:      specLabels,
			}
		}
		return np
	}

	tests := []struct {
		name string
		in   *gkev1beta1.NodePoolList
		want []string
	}{
		{
			name: "keep nodepool without anti keys and name matches current version",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP(currentName, map[string]string{"tidb-component": "shared"}, nil, pointer.String(machineType)),
			}},
			want: []string{currentName},
		},
		{
			name: "keep nodepool without anti keys and name matches desired version",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP(desiredName, map[string]string{"tidb-component": "shared"}, nil, pointer.String(machineType)),
			}},
			want: []string{desiredName},
		},
		{
			name: "keep both current and desired version nodepools when both exist",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP(currentName, map[string]string{"tidb-component": "shared"}, nil, pointer.String(machineType)),
				makeNP(desiredName, map[string]string{"tidb-component": "shared"}, nil, pointer.String(machineType)),
			}},
			want: []string{currentName, desiredName},
		},
		{
			name: "drop when anti key exists in metadata labels",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP(currentName, map[string]string{cloudv1.ProjectTagKey: "p1"}, nil, pointer.String(machineType)),
			}},
			want: []string{},
		},
		{
			name: "drop when machineType is nil (avoid panic)",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP(currentName, nil, nil, nil),
			}},
			want: []string{},
		},
		{
			name: "drop when spec config is nil (avoid panic)",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP(currentName, nil, nil, pointer.String(machineType)),
				// force config nil even if machine type exists in helper input
				{ObjectMeta: metav1.ObjectMeta{Name: currentName}},
			}},
			want: []string{currentName},
		},
		{
			name: "drop when name does not match desired",
			in: &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{
				makeNP("other", nil, nil, pointer.String(machineType)),
			}},
			want: []string{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := filterLegacyNodePoolsByAntiLabels(cluster, tt.in, currentVersion, desiredVersion)
			gotNames := make([]string, 0, len(got.Items))
			for i := range got.Items {
				gotNames = append(gotNames, got.Items[i].Name)
			}
			if diff := cmp.Diff(tt.want, gotNames); diff != "" {
				t.Fatalf("unexpected nodepool names (-want +got):\n%s", diff)
			}
		})
	}
}

func Test_formatGKEVersion(t *testing.T) {
	tests := []struct {
		version string
		want    string
	}{
		{
			version: "1.27.3-gke.100",
			want:    "1.27",
		},
		{
			version: "1.24.14-gke.1200",
			want:    "1.24",
		},
		{
			version: "1.23.17-gke.10000",
			want:    "1.23",
		},
		{
			version: "1.22.17-gke.14100",
			want:    "1.22",
		},
	}
	for _, tt := range tests {
		t.Run(tt.version, func(t *testing.T) {
			if got := formatGKEVersion(tt.version); got != tt.want {
				t.Errorf("formatGKEVersion() = %v, want %v", got, tt.want)
			}
		})
	}
}

func buildTestNodePool(name, instanceType, version, cluster string, labels map[string]string) *gkev1beta1.NodePool {
	npLabels := map[string]string{}
	for k, v := range labels {
		npLabels[k] = v
	}
	npLabels[consts.K8sVersionLabelKey] = formatGKEVersion(version)
	return &gkev1beta1.NodePool{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Spec: gkev1beta1.NodePoolSpec{
			ForProvider: gkev1beta1.NodePoolParameters{
				Cluster: fmt.Sprintf("projects/tidbcloud-prod-000/locations/us-central1/clusters/%s", cluster),
				Config: &gkev1beta1.NodeConfig{
					Labels:      npLabels,
					MachineType: ptr.To(instanceType),
				},
				Version: ptr.To(version),
			},
		},
	}
}

func buildTestCluster(name, k8s string, apply func(*cloudv1.Cluster)) *cloudv1.Cluster {
	c := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: "prod-123456789",
		},
		Spec: cloudv1.ClusterSpec{
			Kubernetes: cloudv1.Kubernetes{
				Name: k8s,
			},
		},
	}
	if apply != nil {
		apply(c)
	}
	return c
}

func buildTestPod(name, namespace, component, nodeName string) *corev1.Pod {
	return &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			Labels: map[string]string{
				consts.ManagedByLabelKey: "tidb-operator",
				consts.ComponentLabelKey: component,
				consts.InstanceLabelKey:  consts.ClusterPrefix,
			},
		},
		Spec: corev1.PodSpec{
			NodeName: nodeName,
		},
	}
}

func buildTestNode(name, nodePoolName string) *corev1.Node {
	return &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
			Labels: map[string]string{
				"cloud.google.com/gke-nodepool": nodePoolName,
			},
		},
	}
}

func Test_getTiDBNodePoolsByCluster(t *testing.T) {
	tests := []struct {
		name           string
		cluster        *cloudv1.Cluster
		desiredVersion string
		nodePoolList   *gkev1beta1.NodePoolList
		makeNodePools  func() (oldNPs, newNPs []*gkev1beta1.NodePool)
		wantNeedCreate bool
		wantErr        bool
	}{
		{
			name:           "old nodepools should not contain other gke's nodepools and other instance type' nodepools",
			desiredVersion: "1.31",
			cluster: buildTestCluster("cluster1", "gke1", func(cluster *cloudv1.Cluster) {
				cluster.Spec.Cluster = cloudv1.ClusterTopology{
					PD:   cloudv1.PD{InstanceType: "n1-standard-4"},
					TiDB: cloudv1.TiDB{InstanceType: "n2-standard-4"},
				}
			}),
			nodePoolList: &gkev1beta1.NodePoolList{
				Items: []gkev1beta1.NodePool{
					// other gke
					*buildTestNodePool("np3", "n2-standard-4", "1.30", "gke2", nil),
					// other instance type
					*buildTestNodePool("np4", "not-used", "1.31", "gke1", nil),
				},
			},
			makeNodePools: func() (oldNPs, newNPs []*gkev1beta1.NodePool) {
				oldNPs = []*gkev1beta1.NodePool{
					buildTestNodePool("np1", "n1-standard-4", "1.30", "gke1", nil),
					buildTestNodePool("np2", "n2-standard-4", "1.30", "gke1", nil),
				}
				newNPs = []*gkev1beta1.NodePool{
					buildTestNodePool("np1-new", "n1-standard-4", "1.31", "gke1", nil),
					buildTestNodePool("np2-new", "n2-standard-4", "1.31", "gke1", nil),
				}
				return oldNPs, newNPs
			},
			wantNeedCreate: false,
		},
		{
			name:           "old nodepools should not contain non-3-layer nodepools",
			desiredVersion: "1.31",
			cluster: buildTestCluster("cluster1", "gke1", func(cluster *cloudv1.Cluster) {
				cluster.Annotations = map[string]string{
					"tidbcloud.com/use-3-layer-node-pool": "true",
				}
				cluster.Spec.Cluster = cloudv1.ClusterTopology{
					PD:   cloudv1.PD{InstanceType: "n1-standard-4"},
					TiDB: cloudv1.TiDB{InstanceType: "n2-standard-4"},
				}
			}),
			nodePoolList: &gkev1beta1.NodePoolList{
				Items: []gkev1beta1.NodePool{
					// does not have 3-layer label
					*buildTestNodePool("np3", "n2-standard-4", "1.30", "gke1", nil),
					*buildTestNodePool("np3-new", "n2-standard-4", "1.31", "gke2", map[string]string{
						"cloud.pingcap.com/for-3-layer-arch": "true",
					}),
				},
			},
			makeNodePools: func() (oldNPs, newNPs []*gkev1beta1.NodePool) {
				labels := map[string]string{
					"cloud.pingcap.com/for-3-layer-arch": "true",
				}
				oldNPs = []*gkev1beta1.NodePool{
					buildTestNodePool("np1", "n1-standard-4", "1.30", "gke1", labels),
				}
				newNPs = []*gkev1beta1.NodePool{
					buildTestNodePool("np1-new", "n1-standard-4", "1.31", "gke1", labels),
				}
				return oldNPs, newNPs
			},
			wantNeedCreate: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			wantOldNPs, wantNewNPs := tt.makeNodePools()
			for _, np := range wantOldNPs {
				tt.nodePoolList.Items = append(tt.nodePoolList.Items, *np)
			}
			for _, np := range wantNewNPs {
				tt.nodePoolList.Items = append(tt.nodePoolList.Items, *np)
			}

			gotOldNPs, gotNewNPs, gotNeedCreate, err := getTiDBNodePoolsByCluster(tt.nodePoolList, tt.cluster, tt.desiredVersion)
			if (err != nil) != tt.wantErr {
				t.Errorf("getTiDBNodePoolsByCluster() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(gotOldNPs, wantOldNPs) {
				t.Errorf("getTiDBNodePoolsByCluster() gotOldNPs = %v, want %v", gotOldNPs, wantOldNPs)
			}
			if !reflect.DeepEqual(gotNewNPs, wantNewNPs) {
				t.Errorf("getTiDBNodePoolsByCluster() gotNewNPs = %v, want %v", gotNewNPs, wantNewNPs)
			}
			if gotNeedCreate != tt.wantNeedCreate {
				t.Errorf("getTiDBNodePoolsByCluster() gotNeedCreate = %v, want %v", gotNeedCreate, tt.wantNeedCreate)
			}
		})
	}
}

func Test_getAllTiDBNodePoolsUnderNetwork(t *testing.T) {
	tests := []struct {
		name string
		gke  *containerpb.Cluster
		want *gkev1beta1.NodePoolList
	}{
		{
			name: "happy path",
			gke: &containerpb.Cluster{
				NodePools: []*containerpb.NodePool{
					{
						Name: "admin-np",
						Config: &containerpb.NodeConfig{
							Labels: map[string]string{
								"cloud.pingcap.com/node-group-type": "control-plane",
								"usedby":                            "dbaas-control-plane-shoot",
							},
						},
					},
					{
						Name: "pd-np",
						Config: &containerpb.NodeConfig{
							Labels: map[string]string{
								"cluster": "shared",
								"usedby":  "customer-tidb",
							},
						},
					},
					{
						Name: "tidb-np",
						Config: &containerpb.NodeConfig{
							Labels: map[string]string{
								"cluster": "shared",
								"usedby":  "customer-tidb",
							},
						},
					},
					{
						Name: "ticdc-np",
						Config: &containerpb.NodeConfig{
							Labels: map[string]string{
								"cluster":   "123456789",
								"component": "ticdc",
								"usedby":    "customer-tidb",
							},
						},
					},
				},
			},
			want: &gkev1beta1.NodePoolList{
				Items: []gkev1beta1.NodePool{
					{ObjectMeta: metav1.ObjectMeta{Name: "pd-np"}},
					{ObjectMeta: metav1.ObjectMeta{Name: "tidb-np"}},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := runtime.NewScheme()
			_ = gkev1beta1.SchemeBuilder.AddToScheme(s)
			infraCli := fake.NewClientBuilder().WithScheme(s).WithLists(tt.want).Build()
			if got := getAllTiDBNodePoolsUnderNetwork(context.TODO(), infraCli, logr.Discard(), tt.gke); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("getAllTiDBNodePoolsUnderNetwork() = %v, want %v", got, tt.want)
			}
		})
	}
}

func makePodTerminating(pod *corev1.Pod) *corev1.Pod {
	now := metav1.Now()
	pod.DeletionTimestamp = &now
	pod.Finalizers = []string{"test.io/finalizer"}
	return pod
}

func Test_countTCPodsOnNewNodePools(t *testing.T) {
	ctx := context.Background()
	ns := "tidb123456789"

	tests := []struct {
		name      string
		pods      []*corev1.Pod
		nodes     []*corev1.Node
		newNPs    []*gkev1beta1.NodePool
		wantNew   int
		wantTotal int
		wantErr   bool
	}{
		{
			name: "all TC pods are on new node pools",
			pods: []*corev1.Pod{
				buildTestPod("db-pd-0", ns, "pd", "node-a"),
				buildTestPod("db-tikv-0", ns, "tikv", "node-b"),
				buildTestPod("db-tidb-0", ns, "tidb", "node-c"),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-a", "np-new-a"),
				buildTestNode("node-b", "np-new-b"),
				buildTestNode("node-c", "np-new-c"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new-a"}},
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new-b"}},
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new-c"}},
			},
			wantNew:   3,
			wantTotal: 3,
		},
		{
			name: "one pod still on old node pool",
			pods: []*corev1.Pod{
				buildTestPod("db-pd-0", ns, "pd", "node-new"),
				buildTestPod("db-tikv-0", ns, "tikv", "node-old"),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-new", "np-new"),
				buildTestNode("node-old", "np-old"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   1,
			wantTotal: 2,
		},
		{
			name: "pod not yet scheduled",
			pods: []*corev1.Pod{
				buildTestPod("db-pd-0", ns, "pd", ""),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 1,
		},
		{
			name: "no active pods because all are terminating",
			pods: []*corev1.Pod{
				makePodTerminating(buildTestPod("db-pd-0", ns, "pd", "node-new")),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-new", "np-new"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 0,
		},
		{
			name: "no TC pods at all",
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 0,
		},
		{
			name:    "no new node pools provided",
			newNPs:  []*gkev1beta1.NodePool{},
			wantErr: true,
		},
		{
			name: "terminating pod on old node pool is skipped, active pod on new node pool",
			pods: []*corev1.Pod{
				makePodTerminating(buildTestPod("db-tikv-old-0", ns, "tikv", "node-old")),
				buildTestPod("db-tikv-0", ns, "tikv", "node-new"),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-old", "np-old"),
				buildTestNode("node-new", "np-new"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   1,
			wantTotal: 1,
		},
		{
			name: "node not found returns error",
			pods: []*corev1.Pod{
				buildTestPod("db-pd-0", ns, "pd", "missing-node"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 1,
			wantErr:   true,
		},
		{
			name: "single pod on the only new node pool",
			pods: []*corev1.Pod{
				buildTestPod("db-pd-0", ns, "pd", "node-a"),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-a", "np-new"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   1,
			wantTotal: 1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := runtime.NewScheme()
			if err := corev1.AddToScheme(s); err != nil {
				t.Fatalf("failed to add corev1 scheme: %v", err)
			}

			objs := make([]client.Object, 0, len(tt.pods)+len(tt.nodes))
			for _, pod := range tt.pods {
				objs = append(objs, pod)
			}
			for _, node := range tt.nodes {
				objs = append(objs, node)
			}

			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()
			gotNew, gotTotal, err := countTCPodsOnNewNodePools(ctx, cli, ns, tt.newNPs)
			if (err != nil) != tt.wantErr {
				t.Fatalf("countTCPodsOnNewNodePools() error = %v, wantErr %v", err, tt.wantErr)
			}
			if gotNew != tt.wantNew {
				t.Fatalf("countTCPodsOnNewNodePools() gotNew = %v, want %v", gotNew, tt.wantNew)
			}
			if gotTotal != tt.wantTotal {
				t.Fatalf("countTCPodsOnNewNodePools() gotTotal = %v, want %v", gotTotal, tt.wantTotal)
			}
		})
	}
}

func Test_countComponentPodsOnNewNodePools(t *testing.T) {
	ctx := context.Background()
	ns := "tidb123456789"

	tests := []struct {
		name      string
		pods      []*corev1.Pod
		nodes     []*corev1.Node
		newNPs    []*gkev1beta1.NodePool
		wantNew   int
		wantTotal int
		wantErr   bool
	}{
		{
			name: "all tiproxy pods are on new node pools",
			pods: []*corev1.Pod{
				buildTestPod("db-tiproxy-0", ns, string(cloudv1.ComponentTiProxy), "node-a"),
				buildTestPod("db-tiproxy-1", ns, string(cloudv1.ComponentTiProxy), "node-b"),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-a", "np-new-a"),
				buildTestNode("node-b", "np-new-b"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new-a"}},
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new-b"}},
			},
			wantNew:   2,
			wantTotal: 2,
		},
		{
			name: "tiproxy pod is still on old node pool",
			pods: []*corev1.Pod{
				buildTestPod("db-tiproxy-0", ns, string(cloudv1.ComponentTiProxy), "node-old"),
			},
			nodes: []*corev1.Node{
				buildTestNode("node-old", "np-old"),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 1,
		},
		{
			name: "tiproxy pod has not been scheduled",
			pods: []*corev1.Pod{
				buildTestPod("db-tiproxy-0", ns, string(cloudv1.ComponentTiProxy), ""),
			},
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 1,
		},
		{
			name: "no active tiproxy pods",
			newNPs: []*gkev1beta1.NodePool{
				{ObjectMeta: metav1.ObjectMeta{Name: "np-new"}},
			},
			wantNew:   0,
			wantTotal: 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := runtime.NewScheme()
			if err := corev1.AddToScheme(s); err != nil {
				t.Fatalf("failed to add corev1 scheme: %v", err)
			}

			objs := make([]client.Object, 0, len(tt.pods)+len(tt.nodes))
			for _, pod := range tt.pods {
				objs = append(objs, pod)
			}
			for _, node := range tt.nodes {
				objs = append(objs, node)
			}

			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(objs...).Build()
			tc := &opv1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      consts.ClusterPrefix,
					Namespace: ns,
				},
			}
			gotNew, gotTotal, err := countComponentPodsOnNewNodePools(ctx, cli, ns, string(cloudv1.ComponentTiProxy), tt.newNPs, tc)
			if (err != nil) != tt.wantErr {
				t.Fatalf("countComponentPodsOnNewNodePools() error = %v, wantErr %v", err, tt.wantErr)
			}
			if gotNew != tt.wantNew {
				t.Fatalf("countComponentPodsOnNewNodePools() gotNew = %v, want %v", gotNew, tt.wantNew)
			}
			if gotTotal != tt.wantTotal {
				t.Fatalf("countComponentPodsOnNewNodePools() gotTotal = %v, want %v", gotTotal, tt.wantTotal)
			}
		})
	}
}