package k8s

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
)

func TestGetComponentTopology(t *testing.T) {
	cases := map[string]struct {
		pods   []*corev1.Pod
		nodes  []*corev1.Node
		expect *ComponentTopology
	}{
		"Normal": {
			pods: []*corev1.Pod{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-0",
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-1",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-1",
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-2",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-2",
				},
			}},
			nodes: []*corev1.Node{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						ZoneLabel: "us-west-2a",
					},
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-1",
					Labels: map[string]string{
						ZoneLabel: "us-west-2b",
					},
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-2",
					Labels: map[string]string{
						ZoneLabel: "us-west-2c",
					},
				},
			}},
			expect: &ComponentTopology{
				Replicas: 3,
				Scheduled: map[string]PodNames{
					"us-west-2a": []string{"tikv-0"},
					"us-west-2b": []string{"tikv-1"},
					"us-west-2c": []string{"tikv-2"},
				},
			},
		},
		"PendingPod": {
			pods: []*corev1.Pod{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-1",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-2",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{},
			}},
			expect: &ComponentTopology{
				Replicas:  3,
				Scheduled: map[string]PodNames{},
				UnScheduled: []string{
					"tikv-0",
					"tikv-1",
					"tikv-2",
				},
			},
		},
		"NoPod": {
			pods: nil,
			expect: &ComponentTopology{
				Replicas:  0,
				Scheduled: map[string]PodNames{},
			},
		},
		"LegacyNodeLabel": {
			pods: []*corev1.Pod{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-0",
				},
			}},
			nodes: []*corev1.Node{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						LegacyZoneLabel: "us-west-2a",
					},
				},
			}},
			expect: &ComponentTopology{
				Replicas: 1,
				Scheduled: map[string]PodNames{
					"us-west-2a": []string{"tikv-0"},
				},
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			ctx := context.TODO()
			cb := fake.NewClientBuilder()
			for i := range tc.pods {
				cb.WithObjects(tc.pods[i])
			}
			for i := range tc.nodes {
				cb.WithObjects(tc.nodes[i])
			}
			client := cb.Build()
			result, err := GetComponentTopology(ctx, client, types.NamespacedName{Namespace: "", Name: "test"}, cloudv1.ComponentTiKV)
			g.Expect(err).To(Succeed())
			g.Expect(result).To(Equal(tc.expect))
		})
	}
}