package eks

import (
	"context"
	"testing"

	"github.com/crossplane-contrib/provider-aws/apis/autoscaling/v1beta1"
	. "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,
				// areAllComponentPodsOnNewNodeGroups filters by InstanceLabelKey.
				// Using the default TidbCluster name ("db") keeps tests aligned with upgrader logic.
				consts.InstanceLabelKey: consts.ClusterPrefix,
			},
		},
		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 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_countTiDBPodsOnNewNodeGroups(t *testing.T) {
	ctx := context.Background()
	ns := "tidb123456789"

	tests := []struct {
		name      string
		pods      []*corev1.Pod
		nodes     []*corev1.Node
		asgs      []*v1beta1.AutoScalingGroup
		wantNew   int
		wantTotal int
		wantMap   map[string][]string
		wantErr   bool
	}{
		{
			name:    "both managed and self managed node groups exist",
			asgs:    getTestAutoScalingGroups(map[string]int64{"test": 1}),
			wantErr: true,
		},
		{
			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,
				},
			),
			wantNew:   3,
			wantTotal: 3,
			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,
				},
			),
			wantNew:   1,
			wantTotal: 3,
			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()
			}

			gotNew, gotTotal, gotMap, err := countTiDBPodsOnNewNodeGroups(ctx, cli, ns, tt.asgs, "")
			if (err != nil) != tt.wantErr {
				t.Errorf("countTiDBPodsOnNewNodeGroups() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if gotNew != tt.wantNew {
				t.Errorf("countTiDBPodsOnNewNodeGroups() gotNew = %v, want %v", gotNew, tt.wantNew)
			}
			if gotTotal != tt.wantTotal {
				t.Errorf("countTiDBPodsOnNewNodeGroups() gotTotal = %v, want %v", gotTotal, tt.wantTotal)
			}
			g.Expect(gotMap).Should(Equal(tt.wantMap))
		})
	}
}

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

	tests := []struct {
		name            string
		pods            []*corev1.Pod
		nodes           []*corev1.Node
		asgs            []*v1beta1.AutoScalingGroup
		extraTiDBNGName string
		wantNew         int
		wantTotal       int
		wantMap         map[string][]string
		wantErr         bool
	}{
		{
			// Bug 2b: 3 new ASGs each have 1 pod, but 2 pods are still on old nodes.
			// Unfixed code: iterates new ASGs, finds each has a pod → returns true (WRONG).
			// Fixed code:   iterates pods, finds 2 pods on old nodes → returns false (CORRECT).
			name: "pods on old nodes with all new ASGs covered (Bug 2b)",
			pods: []*corev1.Pod{
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd-new"),
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb-new"),
				getTestPod("db-tikv-0", ns, string(cloudv1.ComponentTiKV), "node-tikv-new"),
				// These two pods are still on old nodes:
				getTestPod("db-tidb-1", ns, string(cloudv1.ComponentTiDB), "node-tidb-old"),
				getTestPod("db-tikv-1", ns, string(cloudv1.ComponentTiKV), "node-tikv-old"),
			},
			nodes: []*corev1.Node{
				getTestNode("node-pd-new", "ng-pd-new", false),
				getTestNode("node-tidb-new", "ng-tidb-new", false),
				getTestNode("node-tikv-new", "ng-tikv-new", false),
				getTestNode("node-tidb-old", "ng-tidb-old", false),
				getTestNode("node-tikv-old", "ng-tikv-old", false),
			},
			asgs: getTestAutoScalingGroups(map[string]int64{
				"ng-pd-new":   1,
				"ng-tidb-new": 1,
				"ng-tikv-new": 1,
			}),
			wantNew:   3,
			wantTotal: 5,
			wantMap: map[string][]string{
				"ng-pd-new":   {"db-pd-0"},
				"ng-tidb-new": {"db-tidb-0"},
				"ng-tikv-new": {"db-tikv-0"},
				"ng-tidb-old": {"db-tidb-1"},
				"ng-tikv-old": {"db-tikv-1"},
			},
			wantErr: false,
		},
		{
			// Pod on extra tidb NodeGroup (node label "tidb-component: extra-tidb") is accepted
			// as legitimately placed — should NOT cause the function to return false.
			// All other pods are on new ASG nodes.
			name: "pod on extra tidb NodeGroup node is accepted",
			pods: []*corev1.Pod{
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd-new"),
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb-new"),
				getTestPod("db-tikv-0", ns, string(cloudv1.ComponentTiKV), "node-tikv-new"),
				// Extra tidb pod on the extra tidb NodeGroup:
				getTestPod("db-tidb-3", ns, string(cloudv1.ComponentTiDB), "node-extra-tidb"),
			},
			nodes: []*corev1.Node{
				getTestNode("node-pd-new", "ng-pd-new", false),
				getTestNode("node-tidb-new", "ng-tidb-new", false),
				getTestNode("node-tikv-new", "ng-tikv-new", false),
				// Extra tidb NodeGroup node: carries "tidb-component: extra-tidb" label.
				getTestNode("node-extra-tidb", "extra-tidb-ng", true),
			},
			asgs: getTestAutoScalingGroups(map[string]int64{
				"ng-pd-new":   1,
				"ng-tidb-new": 1,
				"ng-tikv-new": 1,
				// "extra-tidb-ng" is NOT in newASGs — it's a manualv1alpha1.NodeGroup, not an ASG.
			}),
			extraTiDBNGName: "extra-tidb-ng",
			wantNew:         4,
			wantTotal:       4,
			wantMap: map[string][]string{
				"ng-pd-new":     {"db-pd-0"},
				"ng-tidb-new":   {"db-tidb-0"},
				"ng-tikv-new":   {"db-tikv-0"},
				"extra-tidb-ng": {"db-tidb-3"},
			},
			wantErr: false,
		},
		{
			// Bug 2a: 2 pods on nodes belonging to the same ASG.
			// Unfixed code: map[string]string overwrites → only 1 pod recorded.
			// Fixed code:   map[string][]string accumulates → both pods recorded.
			name: "multiple pods on same ASG (Bug 2a)",
			pods: []*corev1.Pod{
				getTestPod("db-tidb-0", ns, string(cloudv1.ComponentTiDB), "node-tidb-0"),
				getTestPod("db-tidb-1", ns, string(cloudv1.ComponentTiDB), "node-tidb-1"),
				getTestPod("db-pd-0", ns, string(cloudv1.ComponentPD), "node-pd"),
			},
			nodes: []*corev1.Node{
				// Both tidb nodes belong to the same ASG "ng-tidb-new":
				getTestNode("node-tidb-0", "ng-tidb-new", false),
				getTestNode("node-tidb-1", "ng-tidb-new", false),
				getTestNode("node-pd", "ng-pd-new", false),
			},
			asgs: getTestAutoScalingGroups(map[string]int64{
				"ng-tidb-new": 2,
				"ng-pd-new":   1,
			}),
			wantNew:   3,
			wantTotal: 3,
			wantMap: map[string][]string{
				// Fixed code records BOTH pod names for "ng-tidb-new":
				"ng-tidb-new": {"db-tidb-0", "db-tidb-1"},
				"ng-pd-new":   {"db-pd-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())

			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()

			gotNew, gotTotal, gotMap, err := countTiDBPodsOnNewNodeGroups(ctx, cli, ns, tt.asgs, tt.extraTiDBNGName)
			if (err != nil) != tt.wantErr {
				t.Errorf("countTiDBPodsOnNewNodeGroups() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if gotNew != tt.wantNew {
				t.Errorf("countTiDBPodsOnNewNodeGroups() gotNew = %v, want %v", gotNew, tt.wantNew)
			}
			if gotTotal != tt.wantTotal {
				t.Errorf("countTiDBPodsOnNewNodeGroups() gotTotal = %v, want %v", gotTotal, tt.wantTotal)
			}
			// For the multiple-pods-on-same-ASG case, sort the slice before comparing
			// since pod iteration order is not guaranteed.
			for k := range gotMap {
				if len(gotMap[k]) > 1 {
					sortedGot := make([]string, len(gotMap[k]))
					copy(sortedGot, gotMap[k])
					sortStrings(sortedGot)
					gotMap[k] = sortedGot
				}
			}
			for k := range tt.wantMap {
				if len(tt.wantMap[k]) > 1 {
					sortedWant := make([]string, len(tt.wantMap[k]))
					copy(sortedWant, tt.wantMap[k])
					sortStrings(sortedWant)
					tt.wantMap[k] = sortedWant
				}
			}
			g.Expect(gotMap).Should(Equal(tt.wantMap))
		})
	}
}

// sortStrings sorts a string slice in place (used for deterministic map comparison in tests).
func sortStrings(s []string) {
	for i := 1; i < len(s); i++ {
		for j := i; j > 0 && s[j] < s[j-1]; j-- {
			s[j], s[j-1] = s[j-1], s[j]
		}
	}
}