package mutation

import (
	"context"

	"github.com/pingcap/errors"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	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/metrics"
)

var _ PodMutator = &fastDisk{}

const (
	fastDiskStorageClassName = "fast-disks"
	fastDiskResourceName     = "test-infra.pingcap.net/fast-disks"
)

// fastDisk is a PodMutator that updates pod's fastdisk resource request&limit according it's PVC usage
type fastDisk struct {
	pvcLister  corelisters.PersistentVolumeClaimLister
	kubeClient clientSet.Interface
}

func NewFastDiskPodMutator(kubeClient clientSet.Interface, pvcInformer coreinformers.PersistentVolumeClaimInformer) PodMutator {
	return &fastDisk{
		kubeClient: kubeClient,
		pvcLister:  pvcInformer.Lister(),
	}
}

// Mutate updates pod's fastdisk resource request&limit
func (f *fastDisk) Mutate(ctx context.Context, pod *corev1.Pod) error {
	// get fastdisk usage
	var fastDiskRequest int64
	for _, volume := range pod.Spec.Volumes {
		if volume.PersistentVolumeClaim == nil {
			continue
		}

		pvc, err := f.pvcLister.PersistentVolumeClaims(pod.Namespace).Get(volume.PersistentVolumeClaim.ClaimName)
		switch {
		case apierrors.IsNotFound(err):
			pvc, err = f.kubeClient.CoreV1().PersistentVolumeClaims(pod.Namespace).Get(ctx, volume.PersistentVolumeClaim.ClaimName, metav1.GetOptions{})
			if err != nil {
				return errors.Annotatef(err, "pvc %s/%s not found", pod.Namespace, volume.PersistentVolumeClaim.ClaimName)
			}
			metrics.InformerCacheMissCounter.WithLabelValues("PersistentVolumeClaim").Inc()
		case err != nil:
			return errors.Annotatef(err, "get pvc %s/%s", pod.Namespace, volume.PersistentVolumeClaim.ClaimName)
		}

		if pvc.Spec.StorageClassName != nil && *pvc.Spec.StorageClassName == fastDiskStorageClassName {
			fastDiskRequest++
		}
	}

	if fastDiskRequest > 0 {
		updateFastDiskRequest(pod, fastDiskRequest)
	}

	return nil
}

func updateFastDiskRequest(pod *corev1.Pod, fastDiskRequest int64) {
	// all fast-disk requests are accounted to the first container
	if pod.Spec.Containers[0].Resources.Requests == nil {
		pod.Spec.Containers[0].Resources.Requests = make(corev1.ResourceList)
	}
	pod.Spec.Containers[0].Resources.Requests[fastDiskResourceName] = *resource.NewQuantity(fastDiskRequest, resource.DecimalExponent)

	if pod.Spec.Containers[0].Resources.Limits == nil {
		pod.Spec.Containers[0].Resources.Limits = make(corev1.ResourceList)
	}
	pod.Spec.Containers[0].Resources.Limits[fastDiskResourceName] = *resource.NewQuantity(fastDiskRequest, resource.DecimalExponent)
}