package cloud

import (
	"context"

	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	image = "gcr.io/pingcap-public/third-party/pause:3.9"
)

// Scaler is a helper for Cluster-Autoscaler (CA)
// CA provision new nodes based on pending pods
// During application pods vertical scale or k8s node upgrade,
// the application pods need to be rescheduled to new nodes.
// If we rely on CA to provision node lazily, the service will be
// degraded or even unavailable if there is only one replica.
// With this helper, we can trigger node scale proactively by
// creating placeholder pods
type Scaler interface {
	CheckNodesReady(context.Context, *NodeRequest) (bool, error)
	CreateNodeScaleTrigger(context.Context, *NodeRequest) error
	CleanNodeScaleTrigger(context.Context) error
}

type NodeRequest struct {
	Count        int
	Component    string
	InstanceType string
	NodeVersion  string
	TaintValue   string
}

func getNodeSelector(req *NodeRequest) map[string]string {
	return map[string]string{
		namer.ComponentSelectorKey:  req.Component,
		consts.InstanceTypeLabelKey: req.InstanceType,
		consts.K8sVersionLabelKey:   req.NodeVersion,
	}
}

type scaler struct {
	cli       client.Client
	name      string
	namespace string
}

func NewScaler(cli client.Client, namespace, name string) Scaler {
	return &scaler{
		cli:       cli,
		name:      name,
		namespace: namespace,
	}
}

func (s *scaler) CheckNodesReady(ctx context.Context, req *NodeRequest) (bool, error) {
	podSelector := client.MatchingLabels{
		consts.OverProvisioningComponentLabelKey: req.Component,
		consts.InstanceLabelKey:                  s.name,
	}
	pods := &corev1.PodList{}
	// list pods matching the desired instance type and k8s version
	if err := s.cli.List(ctx, pods, client.InNamespace(s.namespace), podSelector); err != nil {
		return false, err
	}
	scheduled := 0
	desiredSelector := getNodeSelector(req)
	for _, pod := range pods.Items {
		if pod.Spec.NodeName == "" { // pod not scheduled
			continue
		}
		selector := pod.Spec.NodeSelector
		if selector[consts.InstanceTypeLabelKey] != desiredSelector[consts.InstanceTypeLabelKey] {
			continue
		}
		if selector[consts.K8sVersionLabelKey] != desiredSelector[consts.K8sVersionLabelKey] {
			continue
		}
		if !containTaintValue(pod.Spec.Tolerations, req.TaintValue) {
			continue
		}
		scheduled++
	}

	return scheduled == req.Count, nil
}

func (s *scaler) CreateNodeScaleTrigger(ctx context.Context, req *NodeRequest) error {
	labels := map[string]string{
		consts.OverProvisioningComponentLabelKey: req.Component,
		consts.InstanceLabelKey:                  s.name,
	}
	trigger := &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      s.name,
			Namespace: s.namespace,
			Labels:    labels,
		},
		Spec: appsv1.DeploymentSpec{
			Replicas: ptr.To(int32(req.Count)),
			Selector: &metav1.LabelSelector{
				MatchLabels: labels,
			},
			Template: corev1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: labels,
				},
				Spec: corev1.PodSpec{
					TerminationGracePeriodSeconds: ptr.To(int64(0)),
					NodeSelector:                  getNodeSelector(req),
					Affinity: &corev1.Affinity{
						PodAntiAffinity: &corev1.PodAntiAffinity{
							// We use the application namespace and name (as app.kubernetes.io/instance label)
							// to distinguish the overprovision pods
							// So the overprovision pods should not co-schedule to other overprovision pod
							// And node that has other insatnce pod
							RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
								{
									// Avoid scheduling to the node that already has an overprovision pod
									NamespaceSelector: &metav1.LabelSelector{},
									LabelSelector: &metav1.LabelSelector{
										MatchLabels: map[string]string{
											consts.OverProvisioningComponentLabelKey: req.Component,
										},
									},
									TopologyKey: corev1.LabelHostname,
								},
								// Avoid scheduling to the node that has pod belong to other instances
								{ // same namespace but different instance name
									Namespaces: []string{s.namespace},
									LabelSelector: &metav1.LabelSelector{
										MatchExpressions: []metav1.LabelSelectorRequirement{
											{
												Key:      consts.InstanceLabelKey,
												Operator: metav1.LabelSelectorOpNotIn,
												Values:   []string{s.name},
											},
										},
									},
									TopologyKey: corev1.LabelHostname,
								},
								{ // same instance name but different namespace
									NamespaceSelector: &metav1.LabelSelector{
										MatchExpressions: []metav1.LabelSelectorRequirement{
											{
												Key:      corev1.LabelMetadataName,
												Operator: metav1.LabelSelectorOpNotIn,
												Values:   []string{s.namespace},
											},
										},
									},
									LabelSelector: &metav1.LabelSelector{
										MatchExpressions: []metav1.LabelSelectorRequirement{
											{
												Key:      consts.InstanceLabelKey,
												Operator: metav1.LabelSelectorOpIn,
												Values:   []string{s.name},
											},
										},
									},
									TopologyKey: corev1.LabelHostname,
								},
							},
						},
					},
					Tolerations: []corev1.Toleration{
						{
							Key:      "dedicated",
							Operator: corev1.TolerationOpEqual,
							Value:    req.TaintValue,
							Effect:   corev1.TaintEffectNoSchedule,
						},
					},
					Containers: []corev1.Container{
						{
							Name:  "pause",
							Image: image,
							Resources: corev1.ResourceRequirements{
								Requests: corev1.ResourceList{
									corev1.ResourceCPU: resource.MustParse("5m"),
								},
							},
						},
					},
				},
			},
		},
	}
	return k8s.Apply(ctx, s.cli, trigger)
}

func (s *scaler) CleanNodeScaleTrigger(ctx context.Context) error {
	trigger := &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      s.name,
			Namespace: s.namespace,
		},
	}
	return k8s.DeleteObject(ctx, s.cli, trigger)
}

func containTaintValue(tolerations []corev1.Toleration, taintValue string) bool {
	for _, toleration := range tolerations {
		if toleration.Value == taintValue {
			return true
		}
	}
	return false
}