// 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 != ""
// }

package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	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)
	}

	// Load the CA certificate
	caCertFile := os.Getenv("CACERT")
	caCert, err := ioutil.ReadFile(caCertFile)
	if err != nil {
		klog.Fatalf("Error loading CA certificate: %v", err)
	}

	// Create a certificate pool and add the CA certificate
	certPool := x509.NewCertPool()
	if !certPool.AppendCertsFromPEM(caCert) {
		klog.Fatalf("Failed to append CA certificate to the certificate pool")
	}

	// Create TLS credentials with the certificate pool and skip certificate verification
	creds := credentials.NewTLS(&tls.Config{
		InsecureSkipVerify: true,
		RootCAs:            certPool,
	})

	// Create a dial option with the TLS credentials
	dialOption := grpc.WithTransportCredentials(creds)

	deviceConn, err := grpc.Dial("unix:/var/lib/kubelet/device-plugins/mps-device-plugin.sock", dialOption) //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 != ""
}
