package scheduling

// Author: Weisen Pan
// Date: 2023-10-24


import (
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/knets_pkg/api/resource"
	resourcehelper "k8s.io/kubectl/knets_pkg/util/resource"
)

// GreedQueue is used to sort pods using the Greed algorithm
type GreedQueue struct {
	pods          []*v1.Pod
	totalResource v1.ResourceList
}

// NewGreedQueue creates a new GreedQueue instance
func NewGreedQueue(nodes []v1.Node, pods []*v1.Pod) *GreedQueue {
	totalResource := make(v1.ResourceList)
	for _, node := range nodes {
		cpu := totalResource[v1.ResourceCPU]
		memory := totalResource[v1.ResourceMemory]
		cpu.Add(*node.Status.Allocatable.Cpu())
		memory.Add(*node.Status.Allocatable.Memory())
		totalResource[v1.ResourceCPU] = cpu
		totalResource[v1.ResourceMemory] = memory
	}
	return &GreedQueue{
		totalResource: totalResource,
		pods:          pods,
	}
}

// Len returns the number of pods in the queue
func (greed *GreedQueue) Len() int { return len(greed.pods) }

// Swap swaps two pods in the queue
func (greed *GreedQueue) Swap(i, j int) { greed.pods[i], greed.pods[j] = greed.pods[j], greed.pods[i] }

// Less compares two pods based on priority for scheduling
func (greed *GreedQueue) Less(i, j int) bool {
	// DRF algorithm: Prioritize pods with assigned nodes
	if len(greed.pods[i].Spec.NodeName) != 0 {
		return true
	} else if len(greed.pods[j].Spec.NodeName) != 0 {
		return false
	} else {
		return greed.calculatePodShare(greed.pods[i]) > greed.calculatePodShare(greed.pods[j])
	}
}

// calculatePodShare calculates the resource share of a pod
func (greed *GreedQueue) calculatePodShare(pod *v1.Pod) float64 {
	podReq, _ := resourcehelper.PodRequestsAndLimits(pod)
	if len(podReq) == 0 {
		return 0
	}

	var res float64
	for resourceName := range greed.totalResource {
		allocatedRes := podReq[resourceName]
		totalRes := greed.totalResource[resourceName]
		share := Share(allocatedRes.AsApproximateFloat64(), totalRes.AsApproximateFloat64())
		if share > res {
			res = share
		}
	}

	return res
}

// Share calculates the resource share
func Share(alloc, total float64) float64 {
	if total == 0 {
		if alloc == 0 {
			return 0
		}
		return 1
	}
	return alloc / total
}