package validation

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

	"github.com/pingcap/errors"
	admissionv1 "k8s.io/api/admission/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	coreinformers "k8s.io/client-go/informers/core/v1"
	clientSet "k8s.io/client-go/kubernetes"
	corelisters "k8s.io/client-go/listers/core/v1"

	"github.com/pingcap/test-infra/services/webhook/pkg/utils"
)

var _ PodValidator = &holderPodCreation{}

const (
	holderPodLabelKey        = "app.kubernetes.io/component"
	holderPodLabelValue      = "test-infra-holder"
	maxReplicasAnnotationKey = "test-infra.pingcap.net/max-replicas"
	holderPodNamePrefix      = "holder-"
)

// holderPodCreation is a PodValidator that creates holder pod as a side effect of validation
type holderPodCreation struct {
	kubeClient      clientSet.Interface
	podLister       corelisters.PodLister
	namespaceLister corelisters.NamespaceLister
}

func NewHolderPodCreationValidator(kubeClient clientSet.Interface, podInformer coreinformers.PodInformer, namespaceInformer coreinformers.NamespaceInformer) PodValidator {
	return &holderPodCreation{
		kubeClient:      kubeClient,
		podLister:       podInformer.Lister(),
		namespaceLister: namespaceInformer.Lister(),
	}
}

func (h *holderPodCreation) Validate(ctx context.Context, pod *corev1.Pod, opt PodValidationOption) (pass bool, failReason string, err error) {
	if opt.DryRun {
		return true, "", nil
	}

	if !HasPeerHolderPod(pod) {
		return true, "", nil
	}

	if !isStatefulSetPod(pod) {
		return true, "", nil
	}

	switch opt.Operation {
	case admissionv1.Create:
		err = errors.Trace(h.validatePodCreation(ctx, pod))
	case admissionv1.Delete:
		err = errors.Trace(h.validatePodDeletion(ctx, pod))
	}

	if err != nil {
		return false, "internal error", err
	}

	return true, "", nil
}

// validatePodDeletion creates peer holder pod before allowing normal pod to be deleted
func (h *holderPodCreation) validatePodDeletion(ctx context.Context, pod *corev1.Pod) error {
	// skip validation on namespace termination
	ns, err := h.namespaceLister.Get(pod.Namespace)
	if err != nil {
		return errors.Annotatef(err, "get ns of pod %s/%s", pod.Namespace, pod.Name)
	}
	if ns.Status.Phase == corev1.NamespaceTerminating {
		return nil
	}

	holderPodName := BuildHolderPodName(pod.Name)
	_, err = utils.GetPodWithCache(ctx, h.podLister, h.kubeClient, pod.Namespace, holderPodName)
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "get pod %s/%s", pod.Namespace, holderPodName)
	}

	if err == nil { // deleting hook can be called multiple times
		return nil
	}

	holderPod := buildHolderPodFromNormalPod(pod.Name, pod, map[string]string{"kubernetes.io/hostname": pod.Spec.NodeName})
	_, err = h.kubeClient.CoreV1().Pods(pod.Namespace).Create(ctx, holderPod, metav1.CreateOptions{})
	if err != nil {
		return errors.Annotatef(err, "create holderPod %s/%s", pod.Namespace, holderPodName)
	}

	return nil
}

func (h *holderPodCreation) validatePodCreation(ctx context.Context, pod *corev1.Pod) error {
	maxReplicas, err := strconv.Atoi(pod.Annotations[maxReplicasAnnotationKey])
	if err != nil {
		return errors.Annotatef(err, "parse maxReplicas [%s] of pod %s/%s", pod.Annotations[maxReplicasAnnotationKey], pod.Namespace, pod.Name)
	}

	statefulSetName, _ := GetStatefulSetNameAndOrdinal(pod)
	for ordinal := 0; ordinal < maxReplicas; ordinal++ {
		normalPodName := fmt.Sprintf("%s-%d", statefulSetName, ordinal)
		err := h.createPeerHolderPodIfNeeded(ctx, pod.Namespace, normalPodName, pod)
		if err != nil {
			return errors.Trace(err)
		}
	}

	return nil
}

func (h *holderPodCreation) createPeerHolderPodIfNeeded(ctx context.Context, namespace string, normalPodName string, normalPodTemplate *corev1.Pod) error {
	holderPodName := BuildHolderPodName(normalPodName)
	_, err := utils.GetPodWithCache(ctx, h.podLister, h.kubeClient, namespace, holderPodName)
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "get pod %s/%s", namespace, holderPodName)
	}

	if err == nil {
		return nil
	}

	normalPod, err := utils.GetPodWithCache(ctx, h.podLister, h.kubeClient, namespace, normalPodName)
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "get pod %s/%s", namespace, normalPodName)
	}

	if apierrors.IsNotFound(err) ||
		(normalPod.Spec.NodeName == "" && normalPod.Status.NominatedNodeName == "") {
		holderPod := buildHolderPodFromNormalPod(normalPodName, normalPodTemplate, nil)
		_, err = h.kubeClient.CoreV1().Pods(namespace).Create(ctx, holderPod, metav1.CreateOptions{})
		if err != nil {
			return errors.Annotatef(err, "create holderPod %s/%s", namespace, holderPodName)
		}
	}
	return nil
}

func buildHolderPodFromNormalPod(normalPodName string, normalPodTemplate *corev1.Pod, nodeSelector map[string]string) (holderPod *corev1.Pod) {
	holderPodLabels := buildHolderPodLabels(normalPodTemplate)
	holderPod = &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      BuildHolderPodName(normalPodName),
			Namespace: normalPodTemplate.Namespace,
			Labels:    holderPodLabels,
		},
		Spec: corev1.PodSpec{
			TerminationGracePeriodSeconds: &[]int64{0}[0],
			// Spec.Overhead is ignored since we haven't use it
		},
	}

	holderPod.Spec.Tolerations = normalPodTemplate.Spec.Tolerations
	holderPod.Spec.TopologySpreadConstraints = buildHolderPodTopologySpreadConstraints(normalPodTemplate, metav1.LabelSelector{MatchLabels: holderPodLabels})
	holderPod.Spec.InitContainers = buildHolderPodInitContainers(normalPodTemplate)
	holderPod.Spec.Containers = buildHolderPodContainers(normalPodTemplate)

	holderPod.Spec.NodeSelector = buildHolderPodNodeSelector(normalPodTemplate, nodeSelector)

	return holderPod
}

func buildHolderPodLabels(normalPod *corev1.Pod) map[string]string {
	normalPodStatefulSet, _ := GetStatefulSetNameAndOrdinal(normalPod)
	holderPodSet := holderPodNamePrefix + normalPodStatefulSet
	return map[string]string{
		"app.kubernetes.io/instance": holderPodSet,
		holderPodLabelKey:            holderPodLabelValue,
	}
}

func buildHolderPodTopologySpreadConstraints(normalPod *corev1.Pod, labelSelector metav1.LabelSelector) []corev1.TopologySpreadConstraint {
	result := make([]corev1.TopologySpreadConstraint, 0, len(normalPod.Spec.TopologySpreadConstraints))
	for _, topologySpreadConstraint := range normalPod.Spec.TopologySpreadConstraints {
		result = append(result,
			corev1.TopologySpreadConstraint{
				MaxSkew:           topologySpreadConstraint.MaxSkew,
				TopologyKey:       topologySpreadConstraint.TopologyKey,
				WhenUnsatisfiable: topologySpreadConstraint.WhenUnsatisfiable,
				LabelSelector:     &labelSelector,
			},
		)
	}

	return result
}

func buildHolderPodInitContainers(normalPod *corev1.Pod) []corev1.Container {
	return buildHolderContainers(normalPod.Spec.InitContainers, "alpine", []string{"/bin/sh", "-c", "ls"})
}

func buildHolderPodContainers(normalPod *corev1.Pod) []corev1.Container {
	return buildHolderContainers(normalPod.Spec.Containers, "alpine", []string{"tail", "-f", "/dev/null"})
}

func buildHolderContainers(originContainers []corev1.Container, image string, command []string) []corev1.Container {
	result := make([]corev1.Container, 0, len(originContainers))
	for _, container := range originContainers {
		result = append(result,
			corev1.Container{
				Name:      container.Name,
				Image:     image,
				Command:   command,
				Resources: container.Resources,
			},
		)
	}

	return result
}

func buildHolderPodNodeSelector(normalPod *corev1.Pod, extraNodeSelector map[string]string) map[string]string {
	result := make(map[string]string)
	for k, v := range normalPod.Spec.NodeSelector {
		result[k] = v
	}
	for k, v := range extraNodeSelector {
		result[k] = v
	}

	return result
}

// statefulPodRegex is a regular expression that extracts the parent StatefulSet and ordinal from the Name of a Pod
var statefulPodRegex = regexp.MustCompile("(.*)-([0-9]+)$")

// GetStatefulSetNameAndOrdinal gets the name of pod's parent StatefulSet and pod's ordinal as extracted from its Name. If
// the Pod was not created by a StatefulSet, its parent is considered to be empty string, and its ordinal is considered
// to be -1.
func GetStatefulSetNameAndOrdinal(pod *corev1.Pod) (string, int) {
	parent := ""
	ordinal := -1
	subMatches := statefulPodRegex.FindStringSubmatch(pod.Name)
	if len(subMatches) < 3 {
		return parent, ordinal
	}
	parent = subMatches[1]
	if i, err := strconv.ParseInt(subMatches[2], 10, 32); err == nil {
		ordinal = int(i)
	}
	return parent, ordinal
}

func isStatefulSetPod(pod *corev1.Pod) bool {
	for _, owner := range pod.OwnerReferences {
		if owner.Kind == "StatefulSet" {
			return true
		}
	}
	return false
}