package mutation

import (
	"context"

	"github.com/pingcap/errors"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	coreinformers "k8s.io/client-go/informers/core/v1"
	clientSet "k8s.io/client-go/kubernetes"
	corelisters "k8s.io/client-go/listers/core/v1"
)

var _ PodMutator = &podNodeSelector{}

const (
	podNodeSelectorAnnotationKey = "scheduler.alpha.kubernetes.io/node-selector"
)

// podNodeSelector is a PodMutator works like simplified version of podnodeselector admission controller, see: https://kubernetes.io/docs/reference/access-authn-authz/admission-controllers/#podnodeselector and https://github.com/kubernetes/kubernetes/blob/47ad357c52713aa692aff3b6906e16f27231bab1/plugin/pkg/admission/podnodeselector/admission.go#L17
// intended for environments like AWS-EKS that cannot use podnodeselector admission controller
type podNodeSelector struct {
	kubeClient      clientSet.Interface
	namespaceLister corelisters.NamespaceLister
}

func NewPodNodeSelectorMutator(kubeClient clientSet.Interface, namespaceInformer coreinformers.NamespaceInformer) PodMutator {
	return &podNodeSelector{
		kubeClient:      kubeClient,
		namespaceLister: namespaceInformer.Lister(),
	}
}

// Mutate updates pod's nodeSelector base on annotation of pod's namespace
func (s *podNodeSelector) Mutate(_ context.Context, pod *corev1.Pod) error {
	ns, err := s.namespaceLister.Get(pod.Namespace)
	if err != nil {
		return errors.Annotatef(err, "get ns of pod %s/%s", pod.Namespace, pod.Name)
	}

	namespaceNodeSelectorString := ns.Annotations[podNodeSelectorAnnotationKey]
	if len(namespaceNodeSelectorString) == 0 {
		return nil
	}

	namespaceNodeSelector, err := labels.ConvertSelectorToLabelsMap(namespaceNodeSelectorString)
	if err != nil {
		return errors.Annotatef(err, "parse selector of pod %s/%s %s", pod.Namespace, pod.Name, namespaceNodeSelectorString)
	}

	if labels.Conflicts(namespaceNodeSelector, pod.Spec.NodeSelector) {
		return errors.Errorf("pod %s/%s nodeSelector %s conflicts with its namespace node selector %s",
			pod.Namespace, pod.Name, labels.Set(pod.Spec.NodeSelector).String(), namespaceNodeSelector.String())
	}

	pod.Spec.NodeSelector = labels.Merge(namespaceNodeSelector, pod.Spec.NodeSelector)
	return nil
}