/*
Copyright 2018 The Kubernetes Authors.
Copyright 2018-2025 The Volcano Authors.

Modifications made by Volcano authors:
- Added support for hierarchical DRF (HDRF) scheduling with tree-based resource allocation
- Enhanced DRF plugin with namespace-aware resource sharing and metrics integration

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package drf

import (
	"fmt"
	"math"
	"strconv"
	"strings"

	v1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"

	"volcano.sh/volcano/pkg/scheduler/api"
	"volcano.sh/volcano/pkg/scheduler/api/helpers"
	"volcano.sh/volcano/pkg/scheduler/framework"
	"volcano.sh/volcano/pkg/scheduler/metrics"
	"volcano.sh/volcano/pkg/scheduler/plugins/util"
)

// PluginName indicates name of volcano scheduler plugin.
const PluginName = "drf"

var shareDelta = 0.000001

// hierarchicalNode represents the node hierarchy
// and the corresponding weight and drf attribute
type hierarchicalNode struct {
	parent *hierarchicalNode
	attr   *drfAttr
	// If the node is a leaf node,
	// request represents the request of the job.
	request   *api.Resource
	weight    float64
	saturated bool
	hierarchy string
	children  map[string]*hierarchicalNode
}

func (node *hierarchicalNode) Clone(parent *hierarchicalNode) *hierarchicalNode {
	newNode := &hierarchicalNode{
		parent: parent,
		attr: &drfAttr{
			share:            node.attr.share,
			dominantResource: node.attr.dominantResource,
			allocated:        node.attr.allocated.Clone(),
		},
		request:   node.request.Clone(),
		weight:    node.weight,
		saturated: node.saturated,
		hierarchy: node.hierarchy,
		children:  nil,
	}
	if node.children != nil {
		newNode.children = map[string]*hierarchicalNode{}
		for _, child := range node.children {
			newNode.children[child.hierarchy] = child.Clone(newNode)
		}
	}
	return newNode
}

// resourceSaturated returns true if any resource of the job is saturated or the job demands fully allocated resource
func resourceSaturated(allocated *api.Resource,
	jobRequest *api.Resource, demandingResources map[v1.ResourceName]bool) bool {
	for _, rn := range allocated.ResourceNames() {
		if allocated.Get(rn) != 0 && jobRequest.Get(rn) != 0 &&
			allocated.Get(rn) >= jobRequest.Get(rn) {
			return true
		}
		if !demandingResources[rn] && jobRequest.Get(rn) != 0 {
			return true
		}
	}
	return false
}

type drfAttr struct {
	share            float64
	dominantResource string
	allocated        *api.Resource
}

func (attr *drfAttr) String() string {
	return fmt.Sprintf("dominant resource <%s>, dominant share %f, allocated %s",
		attr.dominantResource, attr.share, attr.allocated)
}

type drfPlugin struct {
	totalResource  *api.Resource
	totalAllocated *api.Resource

	// Key is Job ID
	jobAttrs map[api.JobID]*drfAttr

	// map[namespaceName]->attr
	namespaceOpts map[string]*drfAttr

	// hierarchical tree root
	hierarchicalRoot *hierarchicalNode

	// Arguments given for the plugin
	pluginArguments framework.Arguments
}

// New return drf plugin
func New(arguments framework.Arguments) framework.Plugin {
	return &drfPlugin{
		totalResource:  api.EmptyResource(),
		totalAllocated: api.EmptyResource(),
		jobAttrs:       map[api.JobID]*drfAttr{},
		namespaceOpts:  map[string]*drfAttr{},
		hierarchicalRoot: &hierarchicalNode{
			attr:      &drfAttr{allocated: api.EmptyResource()},
			request:   api.EmptyResource(),
			hierarchy: "root",
			weight:    1,
			children:  map[string]*hierarchicalNode{},
		},
		pluginArguments: arguments,
	}
}

func (drf *drfPlugin) Name() string {
	return PluginName
}

// HierarchyEnabled returns if hierarchy is enabled
func (drf *drfPlugin) HierarchyEnabled(ssn *framework.Session) bool {
	for _, tier := range ssn.Tiers {
		for _, plugin := range tier.Plugins {
			if plugin.Name != PluginName {
				continue
			}
			return plugin.EnabledHierarchy != nil && *plugin.EnabledHierarchy
		}
	}
	return false
}

func (drf *drfPlugin) compareQueues(root *hierarchicalNode, lqueue *api.QueueInfo, rqueue *api.QueueInfo) float64 {
	lnode := root
	lpaths := strings.Split(lqueue.Hierarchy, "/")
	rnode := root
	rpaths := strings.Split(rqueue.Hierarchy, "/")
	for i, depth := 0, min(len(lpaths), len(rpaths)); i < depth; i++ {
		// Saturated nodes have minimum priority,
		// so that demanding nodes will be popped first.
		if !lnode.saturated && rnode.saturated {
			return -1
		}
		if lnode.saturated && !rnode.saturated {
			return 1
		}
		if lnode.attr.share/lnode.weight == rnode.attr.share/rnode.weight {
			if i < depth-1 {
				lnode = lnode.children[lpaths[i+1]]
				rnode = rnode.children[rpaths[i+1]]
			}
		} else {
			return lnode.attr.share/lnode.weight - rnode.attr.share/rnode.weight
		}
	}
	return 0
}

func (drf *drfPlugin) OnSessionOpen(ssn *framework.Session) {
	// Prepare scheduling data for this session.
	drf.totalResource.Add(ssn.TotalResource)

	klog.V(4).Infof("Total Allocatable %s", drf.totalResource)

	hierarchyEnabled := drf.HierarchyEnabled(ssn)

	for _, job := range ssn.Jobs {
		attr := &drfAttr{
			allocated: api.EmptyResource(),
		}

		for status, tasks := range job.TaskStatusIndex {
			if api.AllocatedStatus(status) {
				for _, t := range tasks {
					attr.allocated.Add(t.Resreq)
				}
			}
		}

		// Calculate the init share of Job
		drf.updateShare(attr)
		if !ssn.IsJobTerminated(job.UID) {
			metrics.UpdateJobShare(job.Namespace, job.Name, attr.share)
		}

		drf.jobAttrs[job.UID] = attr

		if hierarchyEnabled {
			queue := ssn.Queues[job.Queue]
			drf.totalAllocated.Add(attr.allocated)
			drf.UpdateHierarchicalShare(drf.hierarchicalRoot, drf.totalAllocated, job, attr, queue.Hierarchy, queue.Weights)
		}
	}

	preemptableFn := func(preemptor *api.TaskInfo, preemptees []*api.TaskInfo) ([]*api.TaskInfo, int) {
		var victims []*api.TaskInfo

		addVictim := func(candidate *api.TaskInfo) {
			victims = append(victims, candidate)
		}

		latt := drf.jobAttrs[preemptor.Job]
		if latt == nil {
			klog.Warningf("[drf] Skip preemption: preemptor job <%s> not found in jobAttrs (orphaned task from deleted PodGroup)",
				preemptor.Job)
			return nil, util.Permit
		}
		lalloc := latt.allocated.Clone().Add(preemptor.Resreq)
		_, ls := drf.calculateShare(lalloc, drf.totalResource)

		allocations := map[api.JobID]*api.Resource{}

		for _, preemptee := range preemptees {
			if _, found := allocations[preemptee.Job]; !found {
				ratt := drf.jobAttrs[preemptee.Job]
				if ratt == nil {
					klog.Warningf("[drf] Skip preemptee <%s/%s>: job <%s> not found in jobAttrs (orphaned task from deleted PodGroup)",
						preemptee.Namespace, preemptee.Name, preemptee.Job)
					continue
				}
				allocations[preemptee.Job] = ratt.allocated.Clone()
			}
			ralloc := allocations[preemptee.Job].Sub(preemptee.Resreq)
			_, rs := drf.calculateShare(ralloc, drf.totalResource)

			if ls < rs || math.Abs(ls-rs) <= shareDelta {
				addVictim(preemptee)
			}
		}

		klog.V(4).Infof("[drf] Victims from DRF preemptable plugin are %+v", victims)

		return victims, util.Permit
	}

	ssn.AddPreemptableFn(drf.Name(), preemptableFn)

	if hierarchyEnabled {
		queueOrderFn := func(l interface{}, r interface{}) int {
			lv := l.(*api.QueueInfo)
			rv := r.(*api.QueueInfo)
			ret := drf.compareQueues(drf.hierarchicalRoot, lv, rv)
			if ret < 0 {
				return -1
			}
			if ret > 0 {
				return 1
			}
			return 0
		}
		ssn.AddQueueOrderFn(drf.Name(), queueOrderFn)

		reclaimFn := func(reclaimer *api.TaskInfo, reclaimees []*api.TaskInfo) ([]*api.TaskInfo, int) {
			var victims []*api.TaskInfo
			// clone hdrf tree
			totalAllocated := drf.totalAllocated.Clone()
			root := drf.hierarchicalRoot.Clone(nil)

			//  update reclaimer hdrf
			ljob := ssn.Jobs[reclaimer.Job]
			if ljob == nil {
				klog.Warningf("[drf] Skip reclaim: reclaimer job <%s> not found in session (orphaned task from deleted PodGroup)",
					reclaimer.Job)
				return nil, util.Permit
			}
			lqueue := ssn.Queues[ljob.Queue]
			if lqueue == nil {
				klog.V(4).Infof("[drf] Skip reclaim: reclaimer queue <%s> not found in session",
					ljob.Queue)
				return nil, util.Permit
			}
			ljob = ljob.Clone()
			attr := drf.jobAttrs[ljob.UID]
			if attr == nil {
				klog.V(4).Infof("[drf] Skip reclaim: reclaimer job <%s> not found in jobAttrs",
					ljob.UID)
				return nil, util.Permit
			}
			lattr := &drfAttr{
				allocated: attr.allocated.Clone(),
			}
			lattr.allocated.Add(reclaimer.Resreq)
			totalAllocated.Add(reclaimer.Resreq)
			drf.updateShare(lattr)
			drf.UpdateHierarchicalShare(root, totalAllocated, ljob, lattr, lqueue.Hierarchy, lqueue.Weights)

			for _, preemptee := range reclaimees {
				rjob := ssn.Jobs[preemptee.Job]
				if rjob == nil {
					klog.Warningf("[drf] Skip reclaimee <%s/%s>: job <%s> not found in session (orphaned task from deleted PodGroup)",
						preemptee.Namespace, preemptee.Name, preemptee.Job)
					continue
				}
				rqueue := ssn.Queues[rjob.Queue]
				if rqueue == nil {
					klog.V(4).Infof("[drf] Skip reclaimee <%s/%s>: queue <%s> not found in session",
						preemptee.Namespace, preemptee.Name, rjob.Queue)
					continue
				}

				// update hdrf of reclaimee job
				totalAllocated.Sub(preemptee.Resreq)
				rjob = rjob.Clone()
				attr := drf.jobAttrs[rjob.UID]
				if attr == nil {
					klog.V(4).Infof("[drf] Skip reclaimee <%s/%s>: job <%s> not found in jobAttrs",
						preemptee.Namespace, preemptee.Name, rjob.UID)
					totalAllocated.Add(preemptee.Resreq)
					continue
				}
				rattr := &drfAttr{
					allocated: attr.allocated.Clone(),
				}
				rattr.allocated.Sub(preemptee.Resreq)
				drf.updateShare(rattr)
				drf.UpdateHierarchicalShare(root, totalAllocated, rjob, rattr, rqueue.Hierarchy, rqueue.Weights)

				// compare hdrf of queues
				ret := drf.compareQueues(root, lqueue, rqueue)

				// resume hdrf of reclaimee job
				totalAllocated.Add(preemptee.Resreq)
				rattr.allocated.Add(preemptee.Resreq)
				drf.updateShare(rattr)
				drf.UpdateHierarchicalShare(root, totalAllocated, rjob, rattr, rqueue.Hierarchy, rqueue.Weights)

				if ret < 0 {
					victims = append(victims, preemptee)
				}

				if ret > shareDelta {
					continue
				}
			}

			klog.V(4).Infof("[drf] Victims from HDRF plugin are %+v", victims)

			return victims, util.Permit
		}
		ssn.AddReclaimableFn(drf.Name(), reclaimFn)
	}

	jobOrderFn := func(l interface{}, r interface{}) int {
		lv := l.(*api.JobInfo)
		rv := r.(*api.JobInfo)

		klog.V(4).Infof("DRF JobOrderFn: <%v/%v> share state: %v, <%v/%v> share state: %v",
			lv.Namespace, lv.Name, drf.jobAttrs[lv.UID].share, rv.Namespace, rv.Name, drf.jobAttrs[rv.UID].share)

		if drf.jobAttrs[lv.UID].share == drf.jobAttrs[rv.UID].share {
			return 0
		}

		if drf.jobAttrs[lv.UID].share < drf.jobAttrs[rv.UID].share {
			return -1
		}

		return 1
	}

	ssn.AddJobOrderFn(drf.Name(), jobOrderFn)

	// Register event handlers.
	ssn.AddEventHandler(&framework.EventHandler{
		AllocateFunc: func(event *framework.Event) {
			job := ssn.Jobs[event.Task.Job]
			if job == nil {
				klog.Warningf("[drf] Skip allocate event for task <%s/%s>: job <%s> not found in session (orphaned task from deleted PodGroup)",
					event.Task.Namespace, event.Task.Name, event.Task.Job)
				return
			}
			attr := drf.jobAttrs[event.Task.Job]
			if attr == nil {
				klog.Warningf("[drf] Skip allocate event for task <%s/%s>: job <%s> not found in jobAttrs",
					event.Task.Namespace, event.Task.Name, event.Task.Job)
				return
			}
			attr.allocated.Add(event.Task.Resreq)
			drf.updateShare(attr)
			if !ssn.IsJobTerminated(job.UID) {
				metrics.UpdateJobShare(job.Namespace, job.Name, attr.share)
			}

			nsShare := -1.0
			if hierarchyEnabled {
				queue := ssn.Queues[job.Queue]
				if queue != nil {
					drf.totalAllocated.Add(event.Task.Resreq)
					drf.UpdateHierarchicalShare(drf.hierarchicalRoot, drf.totalAllocated, job, attr, queue.Hierarchy, queue.Weights)
				}
			}

			klog.V(4).Infof("[drf] AllocateFunc: task <%v/%v>, resreq <%v>, share <%v>, namespace share <%v>",
				event.Task.Namespace, event.Task.Name, event.Task.Resreq, attr.share, nsShare)
		},
		DeallocateFunc: func(event *framework.Event) {
			job := ssn.Jobs[event.Task.Job]
			if job == nil {
				klog.Warningf("[drf] Skip deallocate event for task <%s/%s>: job <%s> not found in session (orphaned task from deleted PodGroup)",
					event.Task.Namespace, event.Task.Name, event.Task.Job)
				return
			}
			attr := drf.jobAttrs[event.Task.Job]
			if attr == nil {
				klog.Warningf("[drf] Skip deallocate event for task <%s/%s>: job <%s> not found in jobAttrs",
					event.Task.Namespace, event.Task.Name, event.Task.Job)
				return
			}
			attr.allocated.Sub(event.Task.Resreq)
			drf.updateShare(attr)
			if !ssn.IsJobTerminated(job.UID) {
				metrics.UpdateJobShare(job.Namespace, job.Name, attr.share)
			}

			nsShare := -1.0
			if hierarchyEnabled {
				queue := ssn.Queues[job.Queue]
				if queue != nil {
					drf.totalAllocated.Sub(event.Task.Resreq)
					drf.UpdateHierarchicalShare(drf.hierarchicalRoot, drf.totalAllocated, job, attr, queue.Hierarchy, queue.Weights)
				}
			}

			klog.V(4).Infof("[drf] DeallocateFunc: task <%v/%v>, resreq <%v>, share <%v>, namespace share <%v>",
				event.Task.Namespace, event.Task.Name, event.Task.Resreq, attr.share, nsShare)
		},
	})
}

// build hierarchy if the node does not exist
func (drf *drfPlugin) buildHierarchy(root *hierarchicalNode, job *api.JobInfo, attr *drfAttr,
	hierarchy, hierarchicalWeights string) {
	inode := root
	paths := strings.Split(hierarchy, "/")
	weights := strings.Split(hierarchicalWeights, "/")

	for i := 1; i < len(paths); i++ {
		if child, ok := inode.children[paths[i]]; ok {
			inode = child
		} else {
			fweight, _ := strconv.ParseFloat(weights[i], 64)
			if fweight < 1 {
				fweight = 1
			}
			child = &hierarchicalNode{
				weight:    fweight,
				hierarchy: paths[i],
				request:   api.EmptyResource(),
				attr: &drfAttr{
					allocated: api.EmptyResource(),
				},
				children: make(map[string]*hierarchicalNode),
			}
			klog.V(4).Infof("Node %s added to %s, weight %f",
				child.hierarchy, inode.hierarchy, fweight)
			inode.children[paths[i]] = child
			child.parent = inode
			inode = child
		}
	}

	child := &hierarchicalNode{
		weight:    1,
		attr:      attr,
		hierarchy: string(job.UID),
		request:   job.TotalRequest.Clone(),
		children:  nil,
	}
	inode.children[string(job.UID)] = child
	// update drf attribute bottom up
	klog.V(4).Infof("Job <%s/%s> added to %s, weights %s, attr %v, total request: %s",
		job.Namespace, job.Name, inode.hierarchy, hierarchicalWeights, child.attr, job.TotalRequest)
}

// updateHierarchicalShare updates the node attribute recursively
func (drf *drfPlugin) updateHierarchicalShare(node *hierarchicalNode,
	demandingResources map[v1.ResourceName]bool) {
	if node.children == nil {
		node.saturated = resourceSaturated(node.attr.allocated,
			node.request, demandingResources)
		klog.V(4).Infof("Update hierarchical node %s, share %f, dominant %s, resource %v, saturated: %t",
			node.hierarchy, node.attr.share, node.attr.dominantResource, node.attr.allocated, node.saturated)
	} else {
		var mdr float64 = 1
		// get minimum dominant resource share
		for _, child := range node.children {
			drf.updateHierarchicalShare(child, demandingResources)
			// skip empty child and saturated child
			if child.attr.share != 0 && !child.saturated {
				_, resShare := drf.calculateShare(child.attr.allocated, drf.totalResource)
				if resShare < mdr {
					mdr = resShare
				}
			}
		}

		node.attr.allocated = api.EmptyResource()
		saturated := true
		for _, child := range node.children {
			if !child.saturated {
				saturated = false
			}
			// only consider non-empty children
			if child.attr.share != 0 {
				// saturated child is not scaled
				if child.saturated {
					t := child.attr.allocated
					node.attr.allocated.Add(t)
				} else {
					t := child.attr.allocated.Clone().Multi(mdr / child.attr.share)
					node.attr.allocated.Add(t)
				}
			}
		}
		node.attr.dominantResource, node.attr.share = drf.calculateShare(
			node.attr.allocated, drf.totalResource)
		node.saturated = saturated
		klog.V(4).Infof("Update hierarchical node %s, share %f, dominant resource %s, resource %v, saturated: %t",
			node.hierarchy, node.attr.share, node.attr.dominantResource, node.attr.allocated, node.saturated)
	}
}

func (drf *drfPlugin) UpdateHierarchicalShare(root *hierarchicalNode, totalAllocated *api.Resource, job *api.JobInfo, attr *drfAttr, hierarchy, hierarchicalWeights string) {
	// filter out demanding resources
	demandingResources := map[v1.ResourceName]bool{}
	for _, rn := range drf.totalResource.ResourceNames() {
		if totalAllocated.Get(rn) < drf.totalResource.Get(rn) {
			demandingResources[rn] = true
		}
	}
	drf.buildHierarchy(root, job, attr, hierarchy, hierarchicalWeights)
	drf.updateHierarchicalShare(root, demandingResources)
}

func (drf *drfPlugin) updateShare(attr *drfAttr) {
	attr.dominantResource, attr.share = drf.calculateShare(attr.allocated, drf.totalResource)
}

func (drf *drfPlugin) calculateShare(allocated, totalResource *api.Resource) (string, float64) {
	res := float64(0)
	dominantResource := ""
	for _, rn := range totalResource.ResourceNames() {
		share := helpers.Share(allocated.Get(rn), totalResource.Get(rn))
		if share > res {
			res = share
			dominantResource = string(rn)
		}
	}

	return dominantResource, res
}

func (drf *drfPlugin) OnSessionClose(session *framework.Session) {
	// Clean schedule data.
	drf.totalResource = api.EmptyResource()
	drf.totalAllocated = api.EmptyResource()
	drf.jobAttrs = map[api.JobID]*drfAttr{}
}
