package main

import (
	"context"
	"fmt"
	"os"
	"time"

	"google.golang.org/grpc"
	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/util/workqueue"
	"k8s.io/klog"
	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
)

const (
	pluginName    = "scheduler_ext"
	gpuMemoryKey  = "mps.iitb/gpu-vmemory"
	gpuCoresKey   = "mps.iitb/gpu-vcore"
	resyncPeriod  = 5 * time.Minute
	defaultWeight = 1
)

type MPSDevicePluginScheduler struct {
	clientset    kubernetes.Interface
	queue        workqueue.RateLimitingInterface
	deviceConn   *grpc.ClientConn
	deviceClient pluginapi.DevicePluginClient
}

func main() {
	kubeconfig := os.Getenv("KUBECONFIG")
	config, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
	if err != nil {
		klog.Fatalf("Error building kubeconfig: %v", err)
	}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		klog.Fatalf("Error creating clientset: %v", err)
	}

	deviceConn, err := grpc.Dial("unix:/var/lib/kubelet/device-plugins/gpu-device-plugin.sock", grpc.WithInsecure())
	if err != nil {
		klog.Fatalf("Error connecting to device plugin: %v", err)
	}

	deviceClient := pluginapi.NewDevicePluginClient(deviceConn)

	scheduler := &MPSDevicePluginScheduler{
		clientset:    clientset,
		queue:        workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter()),
		deviceConn:   deviceConn,
		deviceClient: deviceClient,
	}

	stopCh := make(chan struct{})
	defer close(stopCh)

	go scheduler.run(stopCh)

	wait.Until(func() {
		if err := scheduler.syncPods(); err != nil {
			klog.Errorf("Error syncing pods: %v", err)
		}
	}, resyncPeriod, stopCh)

	<-stopCh
}

func (s *MPSDevicePluginScheduler) run(stopCh <-chan struct{}) {
	go wait.Until(s.worker, time.Second, stopCh)
	<-stopCh
}

func (s *MPSDevicePluginScheduler) worker() {
	for s.processNextItem() {
	}
}

func (s *MPSDevicePluginScheduler) processNextItem() bool {
	obj, shutdown := s.queue.Get()
	if shutdown {
		return false
	}
	err := func(obj interface{}) error {
		defer s.queue.Done(obj)
		var key string
		var ok bool
		if key, ok = obj.(string); !ok {
			s.queue.Done(obj)
			klog.Errorf("Expected string in workqueue, but got %#v", obj)
			return nil
		}
		if err := s.syncPod(key); err != nil {
			return fmt.Errorf("error syncing pod '%s': %v", key, err)
		}
		s.queue.Forget(obj)
		klog.Infof("Successfully synced pod '%s'", key)
		return nil
	}(obj)
	if err != nil {
		klog.Error(err)
		return true
	}
	return true
}

func (s *MPSDevicePluginScheduler) syncPod(key string) error {
	namespace, name, err := cache.SplitMetaNamespaceKey(key)
	if err != nil {
		return err
	}

	pod, err := s.clientset.CoreV1().Pods(namespace).Get(context.TODO(), name, v1.GetOptions{})
	if err != nil {
		return err
	}

	if isPodScheduled(pod) {
		return nil
	}

	nodeName, err := s.findBestNode(pod)
	if err != nil {
		return err
	}

	boundPod := &corev1.Pod{
		ObjectMeta: v1.ObjectMeta{
			Namespace: pod.Namespace,
			Name:      pod.Name,
		},
		Spec:   pod.Spec,
		Status: pod.Status,
	}
	boundPod.Spec.NodeName = nodeName
	if _, err := s.clientset.CoreV1().Pods(pod.Namespace).Update(context.TODO(), boundPod, v1.UpdateOptions{}); err != nil {
		return err
	}

	return nil
}

func (s *MPSDevicePluginScheduler) findBestNode(pod *corev1.Pod) (string, error) {
	nodes, err := s.clientset.CoreV1().Nodes().List(context.TODO(), v1.ListOptions{
		LabelSelector: "gpu=true",
	})
	if err != nil {
		return "", err
	}

	var bestNode string
	bestScore := -1

	for _, node := range nodes.Items {
		if !s.isMPSAvailable(node) {
			continue
		}

		score := s.calculateNodeScore(node, pod)
		if score > bestScore {
			bestNode = node.Name
			bestScore = score
		}
	}

	return bestNode, nil
}

func (s *MPSDevicePluginScheduler) isMPSAvailable(node corev1.Node) bool {
	if node.Labels == nil {
		return false
	}

	_, mpsEnabled := node.Labels["mps-enabled"]
	return mpsEnabled
}

func (s *MPSDevicePluginScheduler) calculateNodeScore(node corev1.Node, pod *corev1.Pod) int {
	score := 0

	for _, container := range pod.Spec.Containers {
		if container.Resources.Requests != nil {
			gpuMemoryReq, ok := container.Resources.Requests[corev1.ResourceName(gpuMemoryKey)]
			if !ok {
				continue
			}
			gpuMemoryNode, ok := node.Status.Allocatable[corev1.ResourceName(gpuMemoryKey)]
			if !ok {
				continue
			}
			if gpuMemoryReq.Value() <= gpuMemoryNode.Value() {
				score++
			} else {
				score--
			}
		}

		if container.Resources.Requests != nil {
			gpuCoresReq, ok := container.Resources.Requests[corev1.ResourceName(gpuCoresKey)]
			if !ok {
				continue
			}
			gpuCoresNode, ok := node.Status.Allocatable[corev1.ResourceName(gpuCoresKey)]
			if !ok {
				continue
			}
			if gpuCoresReq.Value() <= gpuCoresNode.Value() {
				score++
			} else {
				score--
			}
		}
	}

	return score
}

func (s *MPSDevicePluginScheduler) syncPods() error {
	pods, err := s.clientset.CoreV1().Pods("").List(context.TODO(), v1.ListOptions{
		FieldSelector: "status.phase!=Succeeded,status.phase!=Failed",
	})
	if err != nil {
		return err
	}

	for _, pod := range pods.Items {
		key, err := cache.MetaNamespaceKeyFunc(&pod)
		if err != nil {
			klog.Errorf("Error getting key for pod '%s/%s': %v", pod.Namespace, pod.Name, err)
			continue
		}
		s.queue.Add(key)
	}

	return nil
}

func isPodScheduled(pod *corev1.Pod) bool {
	return pod.Spec.NodeName != ""
}
