package cluster

import (
	"context"
	"fmt"
	"regexp"
	"strconv"
	"time"

	"github.com/onsi/ginkgo/v2"
	"github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	labelsutil "github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	"github.com/tidbcloud/aws-shared-provider/tests/e2e/common/helper"
	"github.com/tidbcloud/aws-shared-provider/tests/e2e/common/labels"
	"github.com/tidbcloud/aws-shared-provider/tests/e2e/common/wait"
	"github.com/tidbcloud/aws-shared-provider/tests/framework"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/rand"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	tidbResourcesFormat = `limits:
  cpu: "2"
  memory: 512Mi
requests:
  cpu: %dm
  memory: 256Mi
`
)

var _ = ginkgo.Describe("Cluster", func() {
	f := framework.New()

	framework.It(labels.With(
		labels.Feature("tidbpool"),
	), "check tidb pool scale and upgrade", func() {
		var tidbPodList *corev1.PodList
		var oldMinPoolReplicas, oldMaxPoolReplicas int
		var oldRevision string
		ctx := context.Background()
		listOptions := &client.ListOptions{Namespace: f.TiDBPoolNamespace}
		client.MatchingLabels{
			labelsutil.NameLabelKey:      labelsutil.TidbPodNameLabelValue,
			labelsutil.ComponentLabelKey: labelsutil.TidbPodComponentLabelValue,
		}.ApplyToList(listOptions)

		ginkgo.By("fetch serverless-pool-config")
		serverlessPoolConfig := &corev1.ConfigMap{}
		gomega.Expect(f.Client.Get(ctx, client.ObjectKey{
			Name: types.ServerlessPoolConfig, Namespace: f.TiDBPoolNamespace,
		}, serverlessPoolConfig)).NotTo(gomega.HaveOccurred())
		oldMinPoolReplicas, _ = strconv.Atoi(serverlessPoolConfig.Data["min-pool-replicas"])
		oldMaxPoolReplicas, _ = strconv.Atoi(serverlessPoolConfig.Data["max-pool-replicas"])

		ginkgo.By("set min-pool-replicas 0, try tidb assignment")
		serverlessPoolConfig.Data["min-pool-replicas"] = "0"
		serverlessPoolConfig.Data["max-pool-replicas"] = "0"
		gomega.Expect(f.Client.Update(ctx, serverlessPoolConfig)).NotTo(gomega.HaveOccurred())
		restTidbPool(ctx, f, listOptions, 0)
		cluster := helper.CreateDefaultCluster(ctx, f, true)
		waitTiDBReady(ctx, cluster, 10*time.Minute, 10*time.Second)
		gomega.Expect(f.InfraClusterCli.Delete(ctx, cluster)).NotTo(gomega.HaveOccurred())
		tidbPodList = &corev1.PodList{}
		gomega.Expect(wait.WaitForObjectList(ctx, f.Client, &corev1.PodList{}, listOptions, func() bool {
			return len(tidbPodList.Items) == 0
		}, helper.ClusterTimeout)).NotTo(gomega.HaveOccurred())

		ginkgo.By("get old tidb revision")
		oldRevision = getTidbPoolRevision(ctx, f.Client, listOptions)

		ginkgo.By("scale out tidb pool 1X by min-pool-replica")
		serverlessPoolConfig.Data["min-pool-replicas"] = strconv.Itoa(oldMinPoolReplicas * 2)
		serverlessPoolConfig.Data["max-pool-replicas"] = strconv.Itoa(oldMaxPoolReplicas * 2)
		gomega.Expect(f.Client.Update(ctx, serverlessPoolConfig)).NotTo(gomega.HaveOccurred())
		tidbPodList = &corev1.PodList{}
		gomega.Expect(wait.WaitForObjectList(ctx, f.Client, tidbPodList, listOptions, func() bool {
			return len(tidbPodList.Items) >= oldMinPoolReplicas*2
		}, helper.ClusterTimeout)).NotTo(gomega.HaveOccurred())
		gomega.Expect(getTidbPoolRevision(ctx, f.Client, listOptions)).To(gomega.Equal(oldRevision))

		ginkgo.By("set tidb pool to oldMinPoolReplicas")
		serverlessPoolConfig.Data["min-pool-replicas"] = strconv.Itoa(oldMinPoolReplicas)
		serverlessPoolConfig.Data["max-pool-replicas"] = strconv.Itoa(oldMaxPoolReplicas)
		gomega.Expect(f.Client.Update(ctx, serverlessPoolConfig)).NotTo(gomega.HaveOccurred())
		gomega.Expect(wait.WaitForObjectList(ctx, f.Client, tidbPodList, listOptions, func() bool {
			return len(tidbPodList.Items) <= oldMaxPoolReplicas
		}, helper.UpgradeTimeout)).NotTo(gomega.HaveOccurred())

		restTidbPool(ctx, f, listOptions, oldMinPoolReplicas)

		ginkgo.By("scale out tidb pool 1X by cluster replicas")
		cluster = helper.CreateDefaultCluster(ctx, f, true)
		cluster.Spec.TiDB.Replicas = oldMinPoolReplicas
		gomega.Expect(f.InfraClusterCli.Update(ctx, cluster)).NotTo(gomega.HaveOccurred())
		time.Sleep(3 * time.Second)
		waitTiDBReady(ctx, cluster, 10*time.Minute, 10*time.Second)
		// already used a tidb pod when create cluster
		for i := 1; i < cluster.Spec.TiDB.Replicas; i++ {
			gomega.Expect(endpointConnected(ctx, cluster)).NotTo(gomega.HaveOccurred())
		}
		tidbPodList = &corev1.PodList{}
		gomega.Expect(wait.WaitForObjectList(ctx, f.Client, tidbPodList, listOptions, func() bool {
			f.Logger.Info("tidb pods in pool", "size", len(tidbPodList.Items))
			return len(tidbPodList.Items) >= oldMinPoolReplicas*2
		}, helper.ClusterTimeout)).NotTo(gomega.HaveOccurred())

		var oldTidbServerPodTemplate, oldTidbResources, oldServerlessTidbConfigFile string

		ginkgo.By("upgrade tidb pool, replace tidb-server-pod-template")
		tidbServerPodTemplate := &corev1.ConfigMap{}
		gomega.Expect(f.Client.Get(ctx, client.ObjectKey{
			Name: types.TidbServerPodTemplateCm, Namespace: f.TiDBPoolNamespace,
		}, tidbServerPodTemplate)).NotTo(gomega.HaveOccurred())
		oldTidbServerPodTemplate = tidbServerPodTemplate.Data["tidb-server-pod-template"]
		oldTidbResources = serverlessPoolConfig.Data[types.TidbResourcesKey]
		tidbServerPodTemplate.Data["tidb-server-pod-template"] = regexp.MustCompile(`(tcp_keepalive_intvl=)\d+`).
			ReplaceAllString(
				tidbServerPodTemplate.Data["tidb-server-pod-template"],
				fmt.Sprintf("${1}%d", rand.IntnRange(5, 75)),
			)
		gomega.Expect(f.Client.Update(ctx, tidbServerPodTemplate)).NotTo(gomega.HaveOccurred())
		gomega.Expect(waitTidbPoolUpgrade(ctx, f, listOptions, oldRevision)).NotTo(gomega.HaveOccurred())

		ginkgo.By("upgrade tidb pool, set tidb-resource requests.cpu")
		oldRevision1 := getTidbPoolRevision(ctx, f.Client, listOptions)
		serverlessPoolConfig.Data[types.TidbResourcesKey] = fmt.Sprintf(tidbResourcesFormat, rand.IntnRange(5, 50))
		gomega.Expect(f.Client.Update(ctx, serverlessPoolConfig)).NotTo(gomega.HaveOccurred())
		gomega.Expect(waitTidbPoolUpgrade(ctx, f, listOptions, oldRevision1)).NotTo(gomega.HaveOccurred())

		ginkgo.By("upgrade tidb pool, set serverless-tidb-config tidb_stmt_summary_file_max_backups")
		oldRevision2 := getTidbPoolRevision(ctx, f.Client, listOptions)
		serverlessTidbConfig := &corev1.ConfigMap{}
		gomega.Expect(f.Client.Get(ctx, client.ObjectKey{
			Name: "serverless-tidb-config", Namespace: f.TiDBPoolNamespace,
		}, serverlessTidbConfig)).NotTo(gomega.HaveOccurred())
		oldServerlessTidbConfigFile = serverlessTidbConfig.Data["config-file"]
		serverlessTidbConfig.Data["config-file"] = regexp.MustCompile(`(tidb_stmt_summary_file_max_backups\s*=\s*)\d+`).
			ReplaceAllString(
				serverlessTidbConfig.Data["config-file"],
				fmt.Sprintf("${1}%d", rand.IntnRange(1, 50)),
			)
		gomega.Expect(f.Client.Update(ctx, serverlessTidbConfig)).NotTo(gomega.HaveOccurred())
		gomega.Expect(waitTidbPoolUpgrade(ctx, f, listOptions, oldRevision2)).NotTo(gomega.HaveOccurred())

		ginkgo.By("upgrade tidb pool, set back annotations and tidb-resource")
		oldRevision3 := getTidbPoolRevision(ctx, f.Client, listOptions)
		// Get tidbServerPodTemplate again, because it may modify by AnnoTidbScalerLock
		gomega.Expect(f.Client.Get(ctx, client.ObjectKey{
			Name: types.TidbServerPodTemplateCm, Namespace: f.TiDBPoolNamespace,
		}, tidbServerPodTemplate)).NotTo(gomega.HaveOccurred())
		tidbServerPodTemplate.Data["tidb-server-pod-template"] = oldTidbServerPodTemplate
		gomega.Expect(f.Client.Update(ctx, tidbServerPodTemplate)).NotTo(gomega.HaveOccurred())
		serverlessPoolConfig.Data[types.TidbResourcesKey] = oldTidbResources
		gomega.Expect(f.Client.Update(ctx, serverlessPoolConfig)).NotTo(gomega.HaveOccurred())
		serverlessTidbConfig.Data["config-file"] = oldServerlessTidbConfigFile
		gomega.Expect(f.Client.Update(ctx, serverlessTidbConfig)).NotTo(gomega.HaveOccurred())
		gomega.Expect(waitTidbPoolUpgrade(ctx, f, listOptions, oldRevision3)).NotTo(gomega.HaveOccurred())
		gomega.Expect(getTidbPoolRevision(ctx, f.Client, listOptions)).To(gomega.Equal(oldRevision))
	})
})

func restTidbPool(ctx context.Context, f *framework.Framework, listOptions *client.ListOptions, replicas int) {
	ginkgo.By(fmt.Sprintf("rest tidb pool size to %d", replicas))
	gomega.Expect(f.Client.DeleteAllOf(ctx, &corev1.Pod{}, &client.DeleteAllOfOptions{
		ListOptions:   *listOptions,
		DeleteOptions: client.DeleteOptions{GracePeriodSeconds: ptr.To[int64](0)},
	})).NotTo(gomega.HaveOccurred())
	time.Sleep(5 * time.Second)
	tidbPodList := &corev1.PodList{}
	gomega.Expect(wait.WaitForObjectList(ctx, f.Client, tidbPodList, listOptions, func() bool {
		f.Logger.Info("restTidbPool: tidb pods in pool", "size", len(tidbPodList.Items))
		return len(tidbPodList.Items) >= replicas
	}, helper.ClusterTimeout)).NotTo(gomega.HaveOccurred())
}

func getTidbPoolRevision(ctx context.Context, cli client.Client, listOptions *client.ListOptions) string {
	var revision string
	tidbPodList := &corev1.PodList{}
	gomega.Expect(wait.WaitForObjectList(ctx, cli, tidbPodList, listOptions, func() bool {
		if len(tidbPodList.Items) == 0 {
			return false
		}
		for i := range tidbPodList.Items {
			podRevision := tidbPodList.Items[i].Labels[types.TidbServerRevision]
			gomega.Expect(podRevision).NotTo(gomega.BeEmpty())
			if revision == "" {
				revision = podRevision
			} else if revision != podRevision {
				return false
			}
		}
		return true
	}, helper.ClusterTimeout)).NotTo(gomega.HaveOccurred())

	return revision
}

func waitTidbPoolUpgrade(
	ctx context.Context,
	f *framework.Framework,
	listOptions *client.ListOptions,
	oldRevision string,
) error {
	tidbPodList := &corev1.PodList{}
	return wait.WaitForObjectList(ctx, f.Client, tidbPodList, listOptions, func() bool {
		if len(tidbPodList.Items) == 0 {
			f.Logger.Info("tidbPodList is empty")
			return false
		}
		for i := range tidbPodList.Items {
			podRevision := tidbPodList.Items[i].Labels[types.TidbServerRevision]
			gomega.Expect(podRevision).NotTo(gomega.BeEmpty())
			if oldRevision == podRevision {
				f.Logger.Info("found oldRevision in tidbPodList", "oldRevision", oldRevision)
				return false
			}
		}
		return true
	}, helper.UpgradeTimeout)
}