package core

import (
	"context"
	"testing"

	"github.com/tidbcloud/tiflash-manager/pkg/config"
	appsv1 "k8s.io/api/apps/v1"
	autoscalingv2 "k8s.io/api/autoscaling/v2"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8sfake "k8s.io/client-go/kubernetes/fake"
)

func TestCloneStatefulSetForPoolReplacesTemplateRefs(t *testing.T) {
	conf := testConfig()
	templatePool := "main"
	poolName := "analytics"

	template := &appsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      conf.StatefulSetName(templatePool),
			Namespace: conf.Tiflash.Namespace,
			Labels: map[string]string{
				"pool":                               templatePool,
				"statefulset.kubernetes.io/pod-name": conf.StatefulSetName(templatePool),
			},
		},
		Spec: appsv1.StatefulSetSpec{
			ServiceName: conf.PeerServiceName(templatePool),
			Selector: &metav1.LabelSelector{
				MatchLabels: map[string]string{
					"pool": templatePool,
				},
			},
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{
						"pool": templatePool,
						"svc":  conf.PeerServiceName(templatePool),
					},
				},
			},
		},
	}

	cloned := cloneStatefulSetForPool(template, poolName, templatePool, conf)

	if cloned.Name != conf.StatefulSetName(poolName) {
		t.Fatalf("unexpected StatefulSet name: %q", cloned.Name)
	}
	if cloned.Spec.ServiceName != conf.PeerServiceName(poolName) {
		t.Fatalf("unexpected service name: %q", cloned.Spec.ServiceName)
	}
	if got := cloned.Spec.Selector.MatchLabels["pool"]; got != poolName {
		t.Fatalf("selector was not rewritten: %q", got)
	}
	if got := cloned.Spec.Template.Labels["svc"]; got != conf.PeerServiceName(poolName) {
		t.Fatalf("template label was not rewritten: %q", got)
	}
	if got := cloned.Labels[managedPoolLabelKey]; got != poolName {
		t.Fatalf("managed pool label missing: %q", got)
	}
}

func TestCloneStatefulSetForPoolKeepsUnrelatedTemplatePoolStrings(t *testing.T) {
	conf := testConfig()
	templatePool := "main"
	poolName := "analytics"

	template := &appsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name: conf.StatefulSetName(templatePool),
			Labels: map[string]string{
				"pool": templatePool,
			},
			Annotations: map[string]string{
				"note": templatePool,
			},
		},
		Spec: appsv1.StatefulSetSpec{
			ServiceName: conf.PeerServiceName(templatePool),
			Selector: &metav1.LabelSelector{
				MatchLabels: map[string]string{
					"pool": templatePool,
				},
			},
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{
						"pool": templatePool,
					},
					Annotations: map[string]string{
						"note": templatePool,
					},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:    "tiflash",
							Command: []string{"--peer-service=" + conf.PeerServiceName(templatePool)},
							Env: []corev1.EnvVar{
								{Name: "ROLE", Value: templatePool},
								{Name: "PEER_SERVICE", Value: conf.PeerServiceName(templatePool)},
							},
						},
					},
				},
			},
		},
	}

	cloned := cloneStatefulSetForPool(template, poolName, templatePool, conf)

	if got := cloned.Spec.Template.Spec.Containers[0].Env[0].Value; got != templatePool {
		t.Fatalf("unrelated env value should stay %q, got %q", templatePool, got)
	}
	if got := cloned.Spec.Template.Spec.Containers[0].Env[1].Value; got != conf.PeerServiceName(poolName) {
		t.Fatalf("peer service env value was not rewritten: %q", got)
	}
	if got := cloned.Spec.Template.Annotations["note"]; got != templatePool {
		t.Fatalf("annotation should stay %q, got %q", templatePool, got)
	}
}

func TestCloneServiceForPoolRewritesSelectorWithoutTouchingUnrelatedAnnotations(t *testing.T) {
	conf := testConfig()
	templatePool := "main"
	poolName := "analytics"

	template := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name: conf.PeerServiceName(templatePool),
			Labels: map[string]string{
				"pool": templatePool,
			},
			Annotations: map[string]string{
				"note": templatePool,
			},
		},
		Spec: corev1.ServiceSpec{
			Selector: map[string]string{
				"pool":    templatePool,
				"service": conf.PeerServiceName(templatePool),
			},
			ClusterIP: corev1.ClusterIPNone,
		},
	}

	cloned := cloneServiceForPool(template, poolName, templatePool, conf)

	if got := cloned.Spec.Selector["pool"]; got != poolName {
		t.Fatalf("selector pool was not rewritten: %q", got)
	}
	if got := cloned.Spec.Selector["service"]; got != conf.PeerServiceName(poolName) {
		t.Fatalf("selector service was not rewritten: %q", got)
	}
	if got := cloned.Annotations["note"]; got != templatePool {
		t.Fatalf("annotation should stay %q, got %q", templatePool, got)
	}
}

func TestCloneHPAForPoolReplacesScaleTarget(t *testing.T) {
	conf := testConfig()
	templatePool := "main"
	poolName := "analytics"

	template := &autoscalingv2.HorizontalPodAutoscaler{
		ObjectMeta: metav1.ObjectMeta{
			Name: conf.HPAName(templatePool),
		},
		Spec: autoscalingv2.HorizontalPodAutoscalerSpec{
			ScaleTargetRef: autoscalingv2.CrossVersionObjectReference{
				Kind: "StatefulSet",
				Name: conf.StatefulSetName(templatePool),
			},
		},
	}

	cloned := cloneHPAForPool(template, poolName, templatePool, conf)
	if cloned.Name != conf.HPAName(poolName) {
		t.Fatalf("unexpected HPA name: %q", cloned.Name)
	}
	if cloned.Spec.ScaleTargetRef.Name != conf.StatefulSetName(poolName) {
		t.Fatalf("unexpected scale target: %q", cloned.Spec.ScaleTargetRef.Name)
	}
}

func TestReconcileDeletesRemovedManagedPoolResources(t *testing.T) {
	conf := testConfig()
	client := k8sfake.NewSimpleClientset(
		&corev1.Service{
			ObjectMeta: managedObjectMeta(conf.PeerServiceName("analytics"), "analytics"),
		},
		&appsv1.StatefulSet{
			ObjectMeta: managedObjectMeta(conf.StatefulSetName("analytics"), "analytics"),
		},
		&autoscalingv2.HorizontalPodAutoscaler{
			ObjectMeta: managedObjectMeta(conf.HPAName("analytics"), "analytics"),
		},
	)

	reconciler := &PoolReconciler{
		conf:      conf,
		k8sClient: client,
	}

	if err := reconciler.Reconcile(context.Background(), map[string]PoolSpec{
		conf.Tiflash.DefaultPool: {},
	}); err != nil {
		t.Fatalf("reconcile returned error: %v", err)
	}

	if _, err := client.CoreV1().Services(conf.Tiflash.Namespace).Get(context.Background(), conf.PeerServiceName("analytics"), metav1.GetOptions{}); err == nil {
		t.Fatalf("expected removed pool service to be deleted")
	}
	if _, err := client.AppsV1().StatefulSets(conf.Tiflash.Namespace).Get(context.Background(), conf.StatefulSetName("analytics"), metav1.GetOptions{}); err == nil {
		t.Fatalf("expected removed pool statefulset to be deleted")
	}
	if _, err := client.AutoscalingV2().HorizontalPodAutoscalers(conf.Tiflash.Namespace).Get(context.Background(), conf.HPAName("analytics"), metav1.GetOptions{}); err == nil {
		t.Fatalf("expected removed pool hpa to be deleted")
	}
}

func testConfig() *config.Config {
	conf := &config.Config{}
	conf.Tiflash.Namespace = "tidb-serverless"
	conf.Tiflash.DefaultPool = "main"
	conf.Tiflash.TemplatePool = "main"
	conf.Tiflash.StatefulSetNameFormat = "tiflash-{POOL}"
	conf.Tiflash.PeerServiceNameFormat = "tiflash-manager-{POOL}-peer"
	conf.Tiflash.HPANameFormat = "tiflash-{POOL}"
	return conf
}

func managedObjectMeta(name, poolName string) metav1.ObjectMeta {
	return metav1.ObjectMeta{
		Name:      name,
		Namespace: "tidb-serverless",
		Labels: map[string]string{
			managedByLabelKey:   managedByLabelValue,
			managedPoolLabelKey: poolName,
		},
	}
}