package node_controller

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"
	"k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"

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

func TestClusterNodeProtectorReconcile(t *testing.T) {
	cases := map[string]struct {
		node     corev1.Node
		pod      corev1.Pod
		priority string

		expectProtection func(*GomegaWithT, *FakeProtectionEditor, *corev1.Node)
	}{
		"IncreasePriority": {
			node: corev1.Node{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test",
					Labels: map[string]string{
						gardener.NodeGroupSetOwnerKey: "tidb1",
					},
				},
			},
			pod: corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tidb",
					Labels: map[string]string{
						label.NameLabelKey:      "tidb-cluster",
						label.ComponentLabelKey: "tidb",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "test",
				},
			},
			priority: machineHighPriority,
			expectProtection: func(g *GomegaWithT, editor *FakeProtectionEditor, node *corev1.Node) {
				// should not process non-storage node
				_, ok := editor.Records["test"]
				g.Expect(ok).To(BeFalse())
				g.Expect(node.Annotations[protectionKey]).To(Equal(""))
			},
		},
		"ResetPriorityIfNoTiDBPodRunning": {
			node: corev1.Node{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test",
					Labels: map[string]string{
						gardener.NodeGroupSetOwnerKey: "tidb1",
					},
				},
			},
			pod: corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "fluentbit",
				},
				Spec: corev1.PodSpec{
					NodeName: "test",
				},
			},
			priority: machineDefaultPriority,
			expectProtection: func(g *GomegaWithT, editor *FakeProtectionEditor, node *corev1.Node) {
				// should not process non-storage node
				_, ok := editor.Records["test"]
				g.Expect(ok).To(BeFalse())
				g.Expect(node.Annotations[protectionKey]).To(Equal(""))
			},
		},
		"AddNodeProtectionIfTiKVRunning": {
			node: corev1.Node{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv1",
					Labels: map[string]string{
						gardener.NodeGroupSetOwnerKey: "tidb1",
					},
				},
			},
			pod: corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tidb",
					Labels: map[string]string{
						label.NameLabelKey:      "tidb-cluster",
						label.ComponentLabelKey: "tikv",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "test",
				},
			},
			priority: machineHighPriority,
			expectProtection: func(g *GomegaWithT, editor *FakeProtectionEditor, node *corev1.Node) {
				protection, ok := editor.Records["tikv1"]
				g.Expect(ok).To(BeTrue())
				g.Expect(protection).To(BeTrue())
				g.Expect(node.Annotations[protectionKey]).To(Equal(protectionEnabled))
			},
		},
		"RemoveNodeProtectionIfNoStorageRunning": {
			node: corev1.Node{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv2",
					Labels: map[string]string{
						gardener.NodeGroupSetOwnerKey: "tidb1",
					},
					Annotations: map[string]string{
						protectionKey: protectionEnabled,
					},
				},
			},
			pod: corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "fluentbit",
				},
				Spec: corev1.PodSpec{
					NodeName: "test",
				},
			},
			priority: machineDefaultPriority,
			expectProtection: func(g *GomegaWithT, editor *FakeProtectionEditor, node *corev1.Node) {
				protection, ok := editor.Records["tikv2"]
				g.Expect(ok).To(BeTrue())
				g.Expect(protection).To(BeFalse())
				g.Expect(node.Annotations[protectionKey]).To(Equal(protectionDisabled))
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			shootClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(&tc.node, &tc.pod).Build()
			ctx := context.TODO()
			editor := NewFakeProtectionEditor()
			protector := &ClusterNodeProtector{
				InfraClient: nil,
				log:         log.NullLogger{},
				shootClient: shootClient,

				protectionEditor: editor,
			}
			_, err := protector.Reconcile(ctx, ctrl.Request{NamespacedName: types.NamespacedName{Name: tc.node.Name}})
			g.Expect(err).To(Succeed())
			node := &corev1.Node{}
			err = shootClient.Get(ctx, types.NamespacedName{Name: tc.node.Name}, node)
			g.Expect(err).To(Succeed())
			g.Expect(node.Annotations[machinePriorityKey]).To(Equal(tc.priority))
			if tc.expectProtection != nil {
				tc.expectProtection(g, editor, node)
			}
		})
	}
}

func TestParseInstanceID(t *testing.T) {
	cases := map[string]struct {
		id       string
		project  string
		zone     string
		instance string
		err      bool
	}{
		"Basic": {
			id:       "gce://pingcap-gardener/us-west1-a/gke-dev-seed-us-west-gardener-node-po-1cbd90b1-pbf6",
			project:  "pingcap-gardener",
			zone:     "us-west1-a",
			instance: "gke-dev-seed-us-west-gardener-node-po-1cbd90b1-pbf6",
			err:      false,
		},
		"BadPrefix": {
			id:  "bad://a/b/c",
			err: true,
		},
		"BadLevels": {
			id:  "gce://a/b",
			err: true,
		},
		"BadElements": {
			id:  "gce:///",
			err: true,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			project, zone, instance, err := parseInstanceID(tc.id)
			if tc.err {
				g.Expect(err).NotTo(Succeed())
			} else {
				g.Expect(project).To(Equal(tc.project))
				g.Expect(zone).To(Equal(tc.zone))
				g.Expect(instance).To(Equal(tc.instance))
			}
		})
	}
}

type FakeProtectionEditor struct {
	Records map[string]bool
}

func NewFakeProtectionEditor() *FakeProtectionEditor {
	return &FakeProtectionEditor{
		Records: map[string]bool{},
	}
}

func (p *FakeProtectionEditor) EditNodeProtection(node *corev1.Node, protection bool) error {
	p.Records[node.Name] = protection
	return nil
}