/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * openFuyao is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY, FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package plugin

import (
	"time"

	"k8s.io/klog/v2"
	"volcano.sh/volcano/pkg/scheduler/api"
	"volcano.sh/volcano/pkg/scheduler/framework"
	"volcano.sh/volcano/pkg/scheduler/plugins/volcano-xpu-plugin/nodelock"
	"volcano.sh/volcano/pkg/scheduler/plugins/volcano-xpu-plugin/util"
)

const (
	reducedLockRetries = 3
	lockRetryInterval  = 500 * time.Millisecond
)

func IsXPUTask(sJob *SchedulerJob, task *api.TaskInfo) bool {
	if _, ok := sJob.Tasks[task.UID]; !ok {
		return false
	}
	if util.IsXPUName(sJob.Tasks[task.UID].ReqXPUName) {
		return true
	}
	return false
}

func (sh *ScheduleHandler) AllocateXPUForTask(task *api.TaskInfo, ssn *framework.Session) {
	if task == nil || ssn == nil {
		klog.V(util.LogErrorLevel).Infof("AllocateXPUForTask %s.", util.ArgumentError)
		return
	}
	sJob, ok := sh.Jobs[task.Job]
	if !ok {
		klog.V(util.LogDebugLevel).Infof("AllocateXPUForTask %s not req npu.", task.Name)
		return
	}
	if !IsXPUTask(sJob, task) {
		klog.V(util.LogDebugLevel).Infof("AllocateXPUForTask %s no need to set pod annotation.", task.Name)
		return
	}
	if !sJob.JobReadyTag {
		klog.V(util.LogDebugLevel).Infof("AllocateXPUForTask %s not allow allocate npu.", task.Name)
		return
	}
	nodeName := task.NodeName
	node, found := ssn.Nodes[nodeName]
	if !found {
		klog.V(util.LogWarningLevel).Infof("%s AllocateXPUForTask %s not exist.", PluginName, nodeName)
		return
	}

	nodelock.UseClient(ssn.KubeClient())
	if task.Pod != nil {
		var lockErr error
		for attempt := 0; attempt < reducedLockRetries; attempt++ {
			lockErr = nodelock.LockNodeWithPod(nodeName, util.VXPULockName, task.Pod)
			if lockErr == nil {
				break
			}
			klog.V(util.LogWarningLevel).Infof(
				"AllocateXPUForTask node %s is locked, waiting for %s (attempt %d/%d): %v",
				nodeName, lockRetryInterval, attempt+1, reducedLockRetries, lockErr)
			time.Sleep(lockRetryInterval)
		}
		if lockErr != nil {
			klog.V(util.LogErrorLevel).Infof(
				"AllocateXPUForTask lock node %s failed after %d attempts: %v",
				nodeName, reducedLockRetries, lockErr)
			return
		}
	}

	sh.GetAllocatableXPUDevicesByNode(sJob)
	err := sJob.handler.Allocate(sJob, task, node, sh.getNodeXPUDevices(nodeName), sh.Templates)
	if err != nil {
		klog.V(util.LogErrorLevel).Infof("AllocateXPUForTask allocate failed: %s.", util.ArgumentError)
		if task.Pod != nil {
			if releaseErr := nodelock.ReleaseNodeLockWithPod(nodeName, util.VXPULockName, task.Pod); releaseErr != nil {
				klog.V(util.LogErrorLevel).Infof("AllocateXPUForTask release lock failed: %v", releaseErr)
			}
		}
	} else {
		if task.Pod != nil {
			nodelock.ReleaseNodeMemoryLockOnly(nodeName)
		}
	}
}

func (sh *ScheduleHandler) DeallocateXPUForTask(task *api.TaskInfo) {
	if sh == nil || task == nil {
		klog.V(util.LogErrorLevel).Infof("DeallocateXPUForTask failed: %s.", util.ArgumentError)
		return
	}
	if task.Pod != nil && task.NodeName != "" {
		if err := nodelock.ReleaseNodeLockAnnotationOnly(task.NodeName, util.VXPULockName, task.Pod); err != nil {
			klog.V(util.LogErrorLevel).Infof("DeallocateXPUForTask release lock failed: %v, node: %s, pod: %s/%s",
				err, task.NodeName, task.Pod.Namespace, task.Pod.Name)
		}
	}
}