package tidbnodescaler

import (
	"context"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/record"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestTrySetScaleDownDisabledNode(t *testing.T) {
	ctx := context.Background()
	now := time.Now()

	cases := []struct {
		name    string
		initAnn map[string]string
		score   int
		wantAnn map[string]string
	}{
		{
			name:    "Nil node annotations",
			initAnn: nil,
			score:   10,
			wantAnn: map[string]string{
				types.AnnoNodeScaleDownDisabledTimestamp: "not-check-value",
				types.CAScaleDownDisabledKey:             "true",
			},
		},
		{
			name: "Node annotations already have scale down disabled timestamp",
			initAnn: map[string]string{
				types.AnnoNodeScaleDownDisabledTimestamp: now.Format(time.RFC3339),
				types.CAScaleDownDisabledKey:             "true",
			},
			score: 20,
			wantAnn: map[string]string{
				types.AnnoNodeScaleDownDisabledTimestamp: now.Format(time.RFC3339),
				types.CAScaleDownDisabledKey:             "true",
			},
		},
		{
			name: "Node set scale down disabled",
			initAnn: map[string]string{
				types.CAScaleDownDisabledKey: "false",
			},
			score: 30,
			wantAnn: map[string]string{
				types.AnnoNodeScaleDownDisabledTimestamp: "not-check-value",
				types.CAScaleDownDisabledKey:             "true",
			},
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			node := &corev1.Node{
				ObjectMeta: metav1.ObjectMeta{
					Name:        "test-node",
					Annotations: mapCopy(c.initAnn),
				},
			}
			r := newFakeReconciler([]runtime.Object{node})
			r.trySetScaleDownDisabledNode(ctx, node, c.score)

			for k, v := range c.wantAnn {
				if v == "not-check-value" {
					g.Expect(node.Annotations[k]).NotTo(BeEmpty())
					node.Annotations[k] = v
				}
			}
			g.Expect(node.Annotations).To(Equal(c.wantAnn))
		})
	}
}

func TestTryRemoveScaleDownDisabledNode(t *testing.T) {
	g := NewGomegaWithT(t)
	ctx := context.Background()
	now := time.Now()

	cases := []struct {
		name    string
		initAnn map[string]string
		score   int
		wantAnn map[string]string
	}{
		{
			name:    "Node annotations is nil",
			initAnn: nil,
			score:   10,
			wantAnn: nil,
		},
		{
			name: "Node annotations does not have scale down disabled timestamp",
			initAnn: map[string]string{
				types.CAScaleDownDisabledKey: "true",
			},
			score: 20,
			wantAnn: map[string]string{
				types.CAScaleDownDisabledKey: "true",
			},
		},
		{
			name: "Node remove scale down disabled",
			initAnn: map[string]string{
				types.AnnoNodeScaleDownDisabledTimestamp: now.Format(time.RFC3339),
				types.CAScaleDownDisabledKey:             "true",
			},
			score:   40,
			wantAnn: nil,
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			node := &corev1.Node{
				ObjectMeta: metav1.ObjectMeta{
					Name:        "test-node",
					Annotations: mapCopy(c.initAnn),
				},
			}
			r := newFakeReconciler([]runtime.Object{node})
			r.tryRemoveScaleDownDisabledNode(ctx, node, c.score)

			g.Expect(node.Annotations).To(Equal(c.wantAnn))
		})
	}
}

func mapCopy(m map[string]string) map[string]string {
	if m == nil {
		return nil
	}
	res := make(map[string]string, len(m))
	for k, v := range m {
		res[k] = v
	}
	return res
}

func newFakeReconciler(objs []runtime.Object) *Reconciler {
	fakeClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithRuntimeObjects(objs...).
		Build()

	return &Reconciler{
		Client:   fakeClient,
		Recorder: record.NewFakeRecorder(4096),
	}
}