package mixeddeploypd

import (
	"context"
	"fmt"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/equality"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/types"

	"sigs.k8s.io/controller-runtime/pkg/client"
	logf "sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/pingcap/errors"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/utils"
	"github.com/tidbcloud/infra-api/pkg/gitops"
	"github.com/tidbcloud/infra-api/pkg/k8s"
	"github.com/tidbcloud/infra-api/pkg/nodegroup"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
)

var logger = logf.Log.WithName("mixed-deploy-pd-with-tikv")

func ShouldMixedDeployPDWithTiKV(
	topo *cloudv1.ClusterTopology,
	timing string,
	tcNamespace string,
	tcName string,
	provider *cloudv1.Provider,
	createTime time.Time,
	pdMixedDeployConfig *cloudv1.MixedDeployPDWithTiKV) (should bool, err error) {

	tikvInstanceSpec, err := provider.FindInstanceType(topo.TiKV.InstanceType)
	if err != nil {
		return false, errors.Annotate(err, "get tikv instance spec to check should use mixed-deploy-pd")
	}
	return shouldMixedDeployPDWithTiKV(topo.PD.Replicas, topo.TiKV.Replicas, topo.TiDB.Replicas,
		topo.TiFlash.Replicas, timing, tcNamespace, tcName, tikvInstanceSpec, createTime, pdMixedDeployConfig), nil
}

func ShouldUseZeroPDNodes(
	pdCnt int,
	tikvCnt int,
	tidbCnt int,
	tiflashCnt int,
	timing string,
	tc *v1alpha1.TidbCluster,
	provider *cloudv1.Provider,
	labelsGetter func(cloudv1.ComponentType) map[string]string,
	shootClient client.Client,
	cluster *cloudv1.Cluster,
	pdMixedDeployConfig *cloudv1.MixedDeployPDWithTiKV) (should bool, useMixedDeploy bool, err error) {

	if !pdMixedDeployConfig.Enabled {
		return false, false, nil
	}
	if tc.Spec.PD == nil || tc.Spec.PD.Tolerations == nil {
		return false, false, nil
	}
	if !utils.IsMatchedToleration(tc.Spec.PD.Tolerations, cloudv1.ComponentTiKV) {
		return false, false, nil
	}

	tcKey := types.NamespacedName{
		Namespace: tc.Namespace,
		Name:      tc.Name,
	}
	tikvInstanceSpec, err := provider.FindInstanceType(cluster.Spec.Cluster.TiKV.InstanceType)
	if err != nil {
		return false, false, errors.Annotate(err, "get tikv instance spec to check should use mixed-deploy-pd")
	}
	useMixedDeploy = shouldMixedDeployPDWithTiKV(pdCnt, tikvCnt, tidbCnt, tiflashCnt, timing, cluster.Namespace, tcKey.String(),
		tikvInstanceSpec, tc.ObjectMeta.CreationTimestamp.Time, pdMixedDeployConfig)
	if !useMixedDeploy {
		return false, false, nil
	}

	if cluster.Status.Phase == cloudv1.ClusterPhaseCreating {
		logger.Info("using mixed-deploy without extra nodes for new cluster", "cluster", tcKey.String())
		return true, true, nil
	}

	// in getting nodes and pods info below, network error when fetching infos of nodes and pods might cause misjudgement,
	// but it's safe, the worst case is create extra PD nodes, then these nodes will be deleted soon
	if labelsGetter == nil {
		labelsGetter = func(component cloudv1.ComponentType) map[string]string {
			name := nodegroup.NodeGroupNamePrefix(cluster.PrefixedId(), string(component))
			return map[string]string{nodegroup.GardenerWorkerPrefixLabelName: name}
		}
	}
	ctx := context.Background()
	pdNodes, err := k8s.SelectNodes(ctx, shootClient, labelsGetter(cloudv1.ComponentPD))
	if err != nil {
		return false, true, err
	}
	pdNodeCnt := fmt.Sprintf("%v", len(pdNodes))
	runningMixedDeployPDs, runningPDCnt, pdPodCnt, err := getRunningPDPodCnt(ctx, shootClient, tcKey, cloudv1.ComponentTiKV)
	if err != nil {
		return len(pdNodes) == 0, true, err
	}

	if runningMixedDeployPDs == pdCnt {
		var recycleAction string
		if pdPodCnt == pdCnt {
			recycleAction = "no pd nodes need to recycle"
		} else {
			recycleAction = "saftly recycle pd nodes"
		}
		logger.Info("all pd pods are running in tikv nodes, "+recycleAction, "cluster", tcKey.String(), "pd nodes", pdNodeCnt)
		return true, true, nil
	} else {
		var recycleAction string
		if len(pdNodes) == 0 {
			recycleAction = "no pd nodes need to recycle"
		} else {
			recycleAction = "pending to recycle pd nodes"
		}
		logger.Info(fmt.Sprintf("%d/%d(running)/%d(expected) pd pods are running in tikv nodes, %s",
			runningMixedDeployPDs, runningPDCnt, pdCnt, recycleAction), "cluster", tcKey.String(), "pd nodes", pdNodeCnt)
		return len(pdNodes) == 0, true, nil
	}
}

func ChangeTCForMixedDeploy(
	tc *v1alpha1.TidbCluster,
	cluster *cloudv1.Cluster,
	configProvider gitops.ConfigProvider,
	provider *cloudv1.Provider,
	labelsGetter func(cloudv1.ComponentType) map[string]string) error {

	tcKey := types.NamespacedName{
		Namespace: tc.Namespace,
		Name:      tc.Name,
	}
	topo := &cluster.Spec.Cluster

	originalResourcesSetting, err := isOriginalResourcesSetting(tc, configProvider, provider, topo)
	if err != nil {
		return err
	}
	mixedDeployNodeSelector := utils.IsMatchedToleration(tc.Spec.PD.Tolerations, cloudv1.ComponentTiKV)
	if mixedDeployNodeSelector && !originalResourcesSetting {
		// it's already a mixed-deploy tc spec
		return nil
	}

	// straight forward for creating, no multiply steps
	if cluster.Status.Phase == cloudv1.ClusterPhaseCreating {
		tc.Spec.PD.Tolerations = utils.Toleration(cloudv1.ComponentTiKV)
		tc.Spec.PD.NodeSelector = labelsGetter(cloudv1.ComponentTiKV)
		err := reduceResourcesInMixedDeploy(tc, configProvider, provider, topo, true, true)
		if err != nil {
			return err
		}
		logger.Info("using mixed-deploy in creating", "cluster", tcKey.String())
		return nil
	}

	if cluster.Status.Phase != cloudv1.ClusterPhaseRunning {
		logger.Info("waiting for running phase to reduce resources setting", "cluster", tcKey.String())
		return nil
	}

	// step 1: tikv use less resources
	tikvResource, err := configProvider.GetResourceSpec(provider, &topo.TiKV)
	if err != nil {
		return errors.Annotate(err, "get tikv resources spec when reduce resources setting")
	}
	tikvIsOriginalResourcesSetting := equality.Semantic.DeepEqual(tc.Spec.TiKV.ResourceRequirements, tikvResource.Resources)
	if tikvIsOriginalResourcesSetting {
		err = reduceResourcesInMixedDeploy(tc, configProvider, provider, topo, false, true)
		if err != nil {
			return err
		}
		logger.Info("using mixed-deploy tikv resources setting", "cluster", tcKey.String())
		return nil
	}

	// TODO: wait for all tikv resources setting are changed here?

	// step 2: pd use less resources, move pd pods to tikv nodes
	if !mixedDeployNodeSelector {
		tc.Spec.PD.Tolerations = utils.Toleration(cloudv1.ComponentTiKV)
		tc.Spec.PD.NodeSelector = labelsGetter(cloudv1.ComponentTiKV)
		logger.Info("moving pd pods to tikv nodes", "cluster", tcKey.String())
	}
	err = reduceResourcesInMixedDeploy(tc, configProvider, provider, topo, true, false)
	if err != nil {
		return err
	}
	logger.Info("using mixed-deploy pd resources setting", "cluster", tcKey.String())
	return nil
}

func RestoreTCForNonMixedDeploy(
	tc *v1alpha1.TidbCluster,
	cluster *cloudv1.Cluster,
	configProvider gitops.ConfigProvider,
	provider *cloudv1.Provider,
	shootClient client.Client,
	labelsGetter func(cloudv1.ComponentType) map[string]string) (err error) {

	desired := &cluster.Spec.Cluster
	originalResourcesSetting, err := isOriginalResourcesSetting(tc, configProvider, provider, desired)
	if err != nil {
		return err
	}
	normalDeployNodeSelector := utils.IsMatchedToleration(tc.Spec.PD.Tolerations, cloudv1.ComponentPD)
	if normalDeployNodeSelector && originalResourcesSetting {
		// it's a non-mixed-deploy tc spec
		return nil
	}

	tcKey := types.NamespacedName{
		Namespace: tc.Namespace,
		Name:      tc.Name,
	}
	if cluster.Status.Phase != cloudv1.ClusterPhaseRunning {
		logger.Info("waiting for running phase to restore resources setting", "cluster", tcKey.String())
		return nil
	}

	// step 1: restore pd pods
	if !normalDeployNodeSelector {
		tc.Spec.PD.NodeSelector = labelsGetter(cloudv1.ComponentPD)
		tc.Spec.PD.Tolerations = utils.Toleration(cloudv1.ComponentPD)
		logger.Info("restoring pd pods to pd nodes", "cluster", tcKey.String())
		return nil
	}

	// step 2: wait for pd pods are all restored
	runningNonMixedDeployPDs, runningPDs, _, err := getRunningPDPodCnt(context.Background(), shootClient, tcKey, cloudv1.ComponentPD)
	if err != nil {
		return err
	}
	pdCnt := int(tc.Spec.PD.Replicas)
	if runningNonMixedDeployPDs != pdCnt {
		logger.Info(fmt.Sprintf("%d/%d(running)/%d(expected) pd pods are running in pd nodes, pending to restore resources setting",
			runningNonMixedDeployPDs, runningPDs, pdCnt), "cluster", tcKey.String())
		return nil
	}

	// step 3: restore resources setting
	err = restoreResourcesInNonMixedDeploy(tc, configProvider, provider, desired)
	if err != nil {
		return err
	}
	logger.Info("all pd pods are running in pd nodes, restoring resources setting", "cluster", tcKey.String())
	return nil
}

func getRunningPDPodCnt(
	ctx context.Context,
	shootClient client.Client,
	tcKey types.NamespacedName,
	matchToleration cloudv1.ComponentType) (matchedToleration int, running int, total int, err error) {

	pdPodList := &corev1.PodList{}
	err = shootClient.List(
		ctx,
		pdPodList,
		client.MatchingLabels{
			label.ComponentLabelKey: string(cloudv1.ComponentPD),
		},
		client.InNamespace(tcKey.Namespace))
	if err != nil {
		logger.Info("can't get pd pods status when using mixed-deploy", "cluster", tcKey.String())
		return 0, 0, 0, err
	}
	for _, pod := range pdPodList.Items {
		if pod.Status.Phase == corev1.PodRunning {
			running += 1
			if utils.IsMatchedToleration(pod.Spec.Tolerations, matchToleration) {
				matchedToleration += 1
			}
		}
	}
	return matchedToleration, running, len(pdPodList.Items), nil
}

// use mixed-deploy on minimum clusters
func shouldMixedDeployPDWithTiKV(
	pdCnt int,
	tikvCnt int,
	tidbCnt int,
	tiflashCnt int,
	timing string,
	tcNamespace string,
	tcName string,
	tikvInstanceSpec cloudv1.InstanceType,
	createTime time.Time,
	pdMixedDeployConfig *cloudv1.MixedDeployPDWithTiKV) bool {

	affected, reason := pdMixedDeployConfig.ClusterShouldBeAffected(pdCnt, tikvCnt, tidbCnt, tiflashCnt, tikvInstanceSpec, tcNamespace, createTime)
	logger.Info("checking for mixed-deploy when "+timing, "cluster", tcName, "pd", pdCnt, "tikv", tikvCnt, "tidb", tidbCnt,
		"tiflash", tiflashCnt, "created", createTime, "affected", affected, "reason", reason, "result", affected)
	return affected
}

func isOriginalResourcesSetting(
	tc *v1alpha1.TidbCluster,
	configProvider gitops.ConfigProvider,
	provider *cloudv1.Provider,
	desired *cloudv1.ClusterTopology) (bool, error) {

	pdResource, err := configProvider.GetResourceSpec(provider, &desired.PD)
	if err != nil {
		return false, errors.Annotate(err, "get pd resources spec when checking original")
	}
	return equality.Semantic.DeepEqual(tc.Spec.PD.ResourceRequirements, pdResource.Resources), nil
}

func restoreResourcesInNonMixedDeploy(
	tc *v1alpha1.TidbCluster,
	configProvider gitops.ConfigProvider,
	provider *cloudv1.Provider,
	desired *cloudv1.ClusterTopology) (err error) {

	pdResource, err := configProvider.GetResourceSpec(provider, &desired.PD)
	if err != nil {
		return errors.Annotate(err, "get pd resources spec when mixed-deploy-pd => normal")
	}
	tikvResource, err := configProvider.GetResourceSpec(provider, &desired.TiKV)
	if err != nil {
		return errors.Annotate(err, "get tikv resources spec when mixed-deploy-pd => normal")
	}
	tc.Spec.PD.ResourceRequirements = pdResource.Resources
	tc.Spec.TiKV.ResourceRequirements = tikvResource.Resources
	return
}

func reduceResourcesInMixedDeploy(
	tc *v1alpha1.TidbCluster,
	configProvider gitops.ConfigProvider,
	provider *cloudv1.Provider,
	topo *cloudv1.ClusterTopology,
	applyToPD bool,
	applyToTiKV bool) error {

	tikvInstanceSpec, err := provider.FindInstanceType(topo.TiKV.InstanceType)
	if err != nil {
		return err
	}

	// for eks-provider
	if tc.Spec.PD.ResourceRequirements.Requests == nil {
		pdResource, err := configProvider.GetResourceSpec(provider, &topo.PD)
		if err != nil {
			return errors.Annotate(err, "get pd resources spec when normal => mixed-deploy-pd")
		}
		tc.Spec.PD.ResourceRequirements = pdResource.Resources
	}
	if tc.Spec.PD.ResourceRequirements.Limits == nil {
		tc.Spec.PD.ResourceRequirements.Limits = corev1.ResourceList{}
	}
	if tc.Spec.TiKV.ResourceRequirements.Limits == nil {
		tc.Spec.TiKV.ResourceRequirements.Limits = corev1.ResourceList{}
	}

	computeResourcesInMixedDeploy(tc, tikvInstanceSpec, applyToPD, applyToTiKV)
	return nil
}

// limit PD and TiKV resources, the performance is tested on 4c/8c instance types:
// * To make sure PD has needed resouces:
//   - Set PD CPU requests to (instance-resource * 10%)
//   - Set PD Memory requests to (TiKV-allocated / 10)
// * To make sure TiKV's resource not affected by PD:
//   - Set PD CPU limits to (instance-resource * 20%)
//   - Set PD Memory limits to (instance-resource / 10)
// * To make sure requests are satisfied
//   - Set TiKV Memory requests to (origin-requested * 60%), to reserve enough memory for misc service
//   - Set TiKV CPU requests to (origin-requested - PD-requested), to reserve enough CPU for misc service
//   - This will not affect TiKV's performance, because other services in the node all have CPU/memory limits
// - Untouched: TiKV limits
func computeResourcesInMixedDeploy(tc *v1alpha1.TidbCluster, tikvInstanceSpec cloudv1.InstanceType, applyToPD bool, applyToTiKV bool) {
	pdResource := &tc.Spec.PD.ResourceRequirements

	tikvInstanceCPU := float64(tikvInstanceSpec.CPU.MilliValue())
	pdRequestCPU := int64(tikvInstanceCPU * 0.1)
	if applyToPD {
		assignReducedResource(pdResource.Requests, corev1.ResourceCPU, *resource.NewMilliQuantity(pdRequestCPU, resource.DecimalSI))
		assignReducedResource(pdResource.Limits, corev1.ResourceCPU, *resource.NewMilliQuantity(pdRequestCPU*2, resource.DecimalSI))
	}

	tikvInstanceMem := tikvInstanceSpec.Memory.Value()
	if applyToPD {
		assignReducedResource(pdResource.Limits, corev1.ResourceMemory, *resource.NewQuantity(tikvInstanceMem/10, resource.BinarySI))
	}

	tikvResource := &tc.Spec.TiKV.ResourceRequirements
	if tikvResource.Requests == nil {
		if applyToPD {
			assignReducedResource(pdResource.Requests, corev1.ResourceMemory, *resource.NewQuantity(tikvInstanceMem/16, resource.BinarySI))
		}
	} else {
		tikvRequestMem := tikvResource.Requests[corev1.ResourceMemory]
		if applyToPD {
			assignReducedResource(pdResource.Requests, corev1.ResourceMemory, *resource.NewQuantity(tikvRequestMem.Value()/10, resource.BinarySI))
		}
		if applyToTiKV {
			tikvResource.Requests[corev1.ResourceMemory] = *resource.NewQuantity(int64(float64(tikvRequestMem.Value())*0.6), resource.BinarySI)
			tikvRequestCPU := tikvResource.Requests[corev1.ResourceCPU]
			tikvResource.Requests[corev1.ResourceCPU] = *resource.NewMilliQuantity(tikvRequestCPU.MilliValue()-pdRequestCPU, resource.DecimalSI)
		}
	}
}

func assignReducedResource(resourceList corev1.ResourceList, resourceType corev1.ResourceName, quantity resource.Quantity) {
	origin, ok := resourceList[resourceType]
	if !ok || origin.IsZero() || origin.Cmp(quantity) > 0 {
		resourceList[resourceType] = quantity
	}
}