package mutation

import (
	"context"

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

var _ PodMutator = &schedulerNamePodMutator{}

const (
	schedulerNameAnnotationKey = "test-infra.pingcap.net/schedulerName"
)

// schedulerNamePodMutator is a PodMutator that modifies pod's schedulerName by annotation of pod's namespace
type schedulerNamePodMutator struct {
	kubeClient      clientSet.Interface
	namespaceLister corelisters.NamespaceLister
}

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

func (s *schedulerNamePodMutator) 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)
	}

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

	pod.Spec.SchedulerName = schedulerName
	return nil
}