package eks

import (
	"context"
	"testing"

	"github.com/crossplane-contrib/provider-aws/apis/autoscaling/v1beta1"
	"github.com/crossplane-contrib/provider-aws/apis/eks/manualv1alpha1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/pointer"
	"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"
)

func getTestPod(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,
			},
		},
		Spec: corev1.PodSpec{
			NodeName: nodeName,
		},
	}
}

func getTestNode(name, ngName string, managed bool) *corev1.Node {
	labels := map[string]string{}
	if managed {
		labels["eks.amazonaws.com/nodegroup"] = ngName
	} else {
		labels["Name"] = ngName
	}

	return &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
		},
	}
}

func getTestNodeGroups(nameAndSize map[string]int32) []*manualv1alpha1.NodeGroup {
	ngs := make([]*manualv1alpha1.NodeGroup, 0, len(nameAndSize))
	for name, size := range nameAndSize {
		ngs = append(ngs, &manualv1alpha1.NodeGroup{
			ObjectMeta: metav1.ObjectMeta{
				Name: name,
			},
			Status: manualv1alpha1.NodeGroupStatus{
				AtProvider: manualv1alpha1.NodeGroupObservation{
					ScalingConfig: manualv1alpha1.NodeGroupScalingConfigStatus{
						DesiredSize: pointer.Int32(size),
					},
				},
			},
		})
	}
	return ngs
}

func getTestAutoScalingGroups(nameAndSize map[string]int64) []*v1beta1.AutoScalingGroup {
	asgs := make([]*v1beta1.AutoScalingGroup, 0, len(nameAndSize))
	for name, size := range nameAndSize {
		asgs = append(asgs, &v1beta1.AutoScalingGroup{
			ObjectMeta: metav1.ObjectMeta{
				Name: name,
			},
			Spec: v1beta1.AutoScalingGroupSpec{
				ForProvider: v1beta1.AutoScalingGroupParameters{
					DesiredCapacity: pointer.Int64(size),
				},
			},
		})
	}
	return asgs
}

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

	tests := []struct {
		name    string
		pods    []*corev1.Pod
		nodes   []*corev1.Node
		ngs     []*manualv1alpha1.NodeGroup
		asgs    []*v1beta1.AutoScalingGroup
		wantRes bool
		wantMap map[string]string
		wantErr bool
	}{
		{
			name:    "both managed and self managed node groups exist",
			ngs:     getTestNodeGroups(map[string]int32{"test": 1}),
			asgs:    getTestAutoScalingGroups(map[string]int64{"test": 1}),
			wantRes: false,
			wantErr: true,
		},
		{
			name: "all pods are on new managed node groups",
			pods: []*corev1.Pod{
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd"),
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb-a"),
				getTestPod("db-tikv-0", ns, string(cloudv1.ComponentTiKV), "node-tikv"),
				getTestPod("db-tikv-0", "tidb56789", string(cloudv1.ComponentTiKV), "other-cluster-tikv-node"),
			},
			nodes: []*corev1.Node{
				getTestNode("node-pd", "ng-pd", true),
				getTestNode("node-tidb-a", "ng-tidb-a", true),
				getTestNode("node-tikv", "ng-tikv", true),
				getTestNode("other-cluster-tikv-node", "other-cluster-tikv-ng", true),
			},
			ngs: getTestNodeGroups(map[string]int32{
				"ng-pd":     1,
				"ng-tidb-a": 1,
				"ng-tidb-b": 0,
				"ng-tikv":   1,
			}),
			wantRes: true,
			wantMap: map[string]string{
				"ng-pd":     "db-pd-0",
				"ng-tidb-a": "db-tidb-0",
				"ng-tikv":   "db-tikv-0",
			},
			wantErr: false,
		},
		{
			name: "NOT all pods are on new managed node groups",
			pods: []*corev1.Pod{
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd"),
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb"),
				getTestPod("db-tikv-0", ns, string(cloudv1.ComponentTiKV), "node-tikv"),
				getTestPod("db-tikv-0", "tidb56789", string(cloudv1.ComponentTiKV), "other-cluster-tikv-node"),
			},
			nodes: []*corev1.Node{
				getTestNode("node-pd", "old-ng-pd", true),
				getTestNode("node-tidb", "old-ng-tidb", true),
				getTestNode("node-tikv", "ng-tikv", true),
				getTestNode("node-tikv2", "other-cluster-tikv-ng", true),
			},
			ngs: getTestNodeGroups(map[string]int32{
				"ng-pd":   1,
				"ng-tidb": 1,
				"ng-tikv": 1,
			}),
			wantRes: false,
			wantMap: map[string]string{
				"old-ng-pd":   "db-pd-0",
				"old-ng-tidb": "db-tidb-0",
				"ng-tikv":     "db-tikv-0",
			},
			wantErr: false,
		},
		{
			name: "all pods are on new self-managed node groups",
			pods: []*corev1.Pod{
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd"),
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb-a"),
				getTestPod("db-tikv-0", ns, string(cloudv1.ComponentTiKV), "node-tikv"),
				getTestPod("db-tikv-0", "tidb56789", string(cloudv1.ComponentTiKV), "other-cluster-tikv-node"),
			},
			nodes: []*corev1.Node{
				getTestNode("node-pd", "ng-pd", false),
				getTestNode("node-tidb-a", "ng-tidb-a", false),
				getTestNode("node-tikv", "ng-tikv", false),
				getTestNode("other-cluster-tikv-node", "other-cluster-tikv-ng", false),
			},
			asgs: getTestAutoScalingGroups(
				map[string]int64{
					"ng-pd":     1,
					"ng-tidb-a": 1,
					"ng-tidb-b": 0,
					"ng-tikv":   1,
				},
			),
			wantRes: true,
			wantMap: map[string]string{
				"ng-pd":     "db-pd-0",
				"ng-tidb-a": "db-tidb-0",
				"ng-tikv":   "db-tikv-0",
			},
			wantErr: false,
		},
		{
			name: "NOT all pods are on new self-managed node groups",
			pods: []*corev1.Pod{
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd"),
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb"),
				getTestPod("db-tikv-0", ns, string(cloudv1.ComponentTiKV), "node-tikv"),
				getTestPod("db-tikv-0", "tidb56789", string(cloudv1.ComponentTiKV), "other-cluster-tikv-node"),
			},
			nodes: []*corev1.Node{
				getTestNode("node-pd", "old-ng-pd", false),
				getTestNode("node-tidb", "old-ng-tidb", false),
				getTestNode("node-tikv", "ng-tikv", false),
				getTestNode("node-tikv2", "other-cluster-tikv-ng", false),
			},
			asgs: getTestAutoScalingGroups(
				map[string]int64{
					"ng-pd":     1,
					"ng-tidb-a": 1,
					"ng-tidb-b": 0,
					"ng-tikv":   1,
				},
			),
			wantRes: false,
			wantMap: map[string]string{
				"old-ng-pd":   "db-pd-0",
				"old-ng-tidb": "db-tidb-0",
				"ng-tikv":     "db-tikv-0",
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			s := runtime.NewScheme()
			g.Expect(corev1.AddToScheme(s)).Should(Succeed())

			var cli client.Client
			if len(tt.pods)+len(tt.nodes) > 0 {
				existObj := make([]client.Object, 0, len(tt.pods)+len(tt.nodes))
				for i := range tt.pods {
					existObj = append(existObj, tt.pods[i])
				}
				for i := range tt.nodes {
					existObj = append(existObj, tt.nodes[i])
				}
				cli = fake.NewClientBuilder().WithScheme(s).WithObjects(existObj...).Build()
			} else {
				cli = fake.NewClientBuilder().WithScheme(s).Build()
			}

			gotRes, gotMap, err := areAllTiDBPodsOnNewNodeGroups(ctx, cli, ns, tt.ngs, tt.asgs)
			if (err != nil) != tt.wantErr {
				t.Errorf("areAllTiDBPodsOnNewNodeGroups() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if gotRes != tt.wantRes {
				t.Errorf("areAllTiDBPodsOnNewNodeGroups() gotRes = %v, want %v", gotRes, tt.wantRes)
			}
			g.Expect(gotMap).Should(Equal(tt.wantMap))
		})
	}
}

func Test_allTiDBNodesUpgraded(t *testing.T) {
	tests := []struct {
		name string
		info *upgradeInfo
		ngs  []*manualv1alpha1.NodeGroup
		asgs []*v1beta1.AutoScalingGroup
		want bool
	}{
		{
			name: "suspend cluster with self-managed nodegroups",
			info: &upgradeInfo{
				cluster:          &cloudv1.Cluster{Status: cloudv1.ClusterStatus{Phase: cloudv1.ClusterPhaseSuspend}},
				useSelfManaged:   true,
				migrateNodeGroup: false,
			},
			asgs: []*v1beta1.AutoScalingGroup{{}},
			want: true,
		},
		{
			name: "try to migrate suspend cluster with managed nodegroups",
			info: &upgradeInfo{
				cluster:           &cloudv1.Cluster{Status: cloudv1.ClusterStatus{Phase: cloudv1.ClusterPhaseSuspend}},
				useSelfManaged:    false,
				migrateNodeGroup:  true,
				desiredK8sVersion: "1.30",
			},
			ngs: []*manualv1alpha1.NodeGroup{
				{
					Status: manualv1alpha1.NodeGroupStatus{
						ResourceStatus: xpv1.ResourceStatus{
							ConditionedStatus: xpv1.ConditionedStatus{
								Conditions: []xpv1.Condition{
									{Type: xpv1.TypeSynced, Status: corev1.ConditionTrue},
									{Type: xpv1.TypeReady, Status: corev1.ConditionTrue},
								},
							},
						},
						AtProvider: manualv1alpha1.NodeGroupObservation{
							Version: "1.30",
							Status:  manualv1alpha1.NodeGroupStatusActive,
						},
					},
				},
			},
			want: true,
		},
		{
			name: "suspend cluster with managed nodegroups",
			info: &upgradeInfo{
				cluster:           &cloudv1.Cluster{Status: cloudv1.ClusterStatus{Phase: cloudv1.ClusterPhaseSuspend}},
				useSelfManaged:    false,
				migrateNodeGroup:  false,
				desiredK8sVersion: "1.30",
			},
			ngs: []*manualv1alpha1.NodeGroup{
				{
					Status: manualv1alpha1.NodeGroupStatus{
						ResourceStatus: xpv1.ResourceStatus{
							ConditionedStatus: xpv1.ConditionedStatus{
								Conditions: []xpv1.Condition{
									{Type: xpv1.TypeSynced, Status: corev1.ConditionTrue},
									{Type: xpv1.TypeReady, Status: corev1.ConditionTrue},
								},
							},
						},
						AtProvider: manualv1alpha1.NodeGroupObservation{
							Version: "1.30",
							Status:  manualv1alpha1.NodeGroupStatusActive,
						},
					},
				},
			},
			want: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := allTiDBNodesUpgraded(tt.info, tt.ngs, tt.asgs); got != tt.want {
				t.Errorf("allTiDBNodesUpgraded() = %v, want %v", got, tt.want)
			}
		})
	}
}