/*
Copyright 2023 The Volcano Authors.

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 gpushare

import (
	"fmt"

	"github.com/pkg/errors"
	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"

	"volcano.sh/volcano/pkg/scheduler/api/devices"
	"volcano.sh/volcano/pkg/scheduler/plugins/util/nodelock"
)

// GPUDevice include gpu id, memory and the pods that are sharing it.
type GPUDevice struct {
	// GPU ID
	ID int
	// The pods that are sharing this GPU
	PodMap map[string]*v1.Pod
	// memory per card
	Memory uint
}

type GPUDevices struct {
	Name string

	Device map[int]*GPUDevice
}

// NewGPUDevice creates a device
func NewGPUDevice(id int, mem uint) *GPUDevice {
	return &GPUDevice{
		ID:     id,
		Memory: mem,
		PodMap: map[string]*v1.Pod{},
	}
}

func NewGPUDevices(name string, node *v1.Node) *GPUDevices {
	if node == nil {
		return nil
	}
	memory, ok := node.Status.Capacity[VolcanoGPUResource]
	if !ok {
		return nil
	}
	totalMemory := memory.Value()

	res, ok := node.Status.Capacity[VolcanoGPUNumber]
	if !ok {
		return nil
	}
	gpuNumber := res.Value()
	if gpuNumber == 0 {
		klog.Warningf("invalid %s=%s", VolcanoGPUNumber, res.String())
		return nil
	}

	memoryPerCard := uint(totalMemory / gpuNumber)
	gpudevices := GPUDevices{}
	gpudevices.Device = make(map[int]*GPUDevice)
	gpudevices.Name = name
	for i := 0; i < int(gpuNumber); i++ {
		gpudevices.Device[i] = NewGPUDevice(i, memoryPerCard)
	}
	unhealthyGPUs := getUnhealthyGPUs(&gpudevices, node)
	for i := range unhealthyGPUs {
		klog.V(4).Infof("delete unhealthy gpu id %d from GPUDevices", unhealthyGPUs[i])
		delete(gpudevices.Device, unhealthyGPUs[i])
	}
	return &gpudevices
}

// GetIgnoredDevices return device names which wish vc-scheduler to ignore
func (gs *GPUDevices) GetIgnoredDevices() []string {
	return []string{""}
}

// AddResource adds the pod to GPU pool if it is assigned
func (gs *GPUDevices) AddResource(pod *v1.Pod) {
	gpuRes := getGPUMemoryOfPod(pod)
	gpuNumRes := getGPUNumberOfPod(pod)
	if gpuRes > 0 || gpuNumRes > 0 {
		ids := GetGPUIndex(pod)
		for _, id := range ids {
			if dev := gs.Device[id]; dev != nil {
				dev.PodMap[string(pod.UID)] = pod
			}
		}
	}
}

// SubResource frees the gpu hold by the pod
func (gs *GPUDevices) SubResource(pod *v1.Pod) {
	gpuRes := getGPUMemoryOfPod(pod)
	gpuNumRes := getGPUNumberOfPod(pod)
	if gpuRes > 0 || gpuNumRes > 0 {
		gs.removePodFromDeviceMap(pod, GetGPUIndex(pod))
	}
}

func (gs *GPUDevices) removePodFromDeviceMap(pod *v1.Pod, ids []int) {
	if pod == nil {
		return
	}
	if len(ids) == 0 {
		for _, dev := range gs.Device {
			delete(dev.PodMap, string(pod.UID))
		}
		return
	}
	for _, id := range ids {
		if dev := gs.Device[id]; dev != nil {
			delete(dev.PodMap, string(pod.UID))
		}
	}
}

func (gs *GPUDevices) HasDeviceRequest(pod *v1.Pod) bool {
	if GpuSharingEnable && getGPUMemoryOfPod(pod) > 0 ||
		GpuNumberEnable && getGPUNumberOfPod(pod) > 0 {
		return true
	}
	return false
}

func (gs *GPUDevices) AddQueueResource(pod *v1.Pod) map[string]float64 {
	return map[string]float64{}
}

func (gs *GPUDevices) Release(kubeClient kubernetes.Interface, pod *v1.Pod) (*devices.DeviceReservation, error) {
	if pod == nil {
		return nil, nil
	}
	ids := GetGPUIndex(pod)
	gs.removePodFromDeviceMap(pod, ids)

	klog.V(4).Infof("predicates with gpu sharing, update pod %s/%s deallocate from node [%s]", pod.Namespace, pod.Name, gs.Name)
	return &devices.DeviceReservation{
		DeviceType:  DeviceName,
		Annotations: devices.AnnotationKeyMap([]string{PredicateTime, GPUIndex}),
	}, nil
}

func (gs *GPUDevices) FilterNode(pod *v1.Pod, schedulePolicy string) (int, string, error) {
	klog.V(4).Infoln("DeviceSharing:Into FitInPod", pod.Name)
	if GpuSharingEnable {
		fit, err := checkNodeGPUSharingPredicate(pod, gs)
		if err != nil || !fit {
			klog.Errorln("deviceSharing err=", err.Error())
			return devices.Unschedulable, fmt.Sprintf("GpuShare %s", err.Error()), err
		}
	}
	if GpuNumberEnable {
		fit, err := checkNodeGPUNumberPredicate(pod, gs)
		if err != nil || !fit {
			klog.Errorln("deviceSharing err=", err.Error())
			return devices.Unschedulable, fmt.Sprintf("GpuNumber %s", err.Error()), err
		}
	}
	klog.V(4).Infoln("DeviceSharing:FitInPod successed")
	return devices.Success, "", nil
}

func (gs *GPUDevices) GetStatus() string {
	return ""
}

// DeepCopy returns a deep copy of GPUDevices for use in dry-run simulation.
func (gs *GPUDevices) DeepCopy() interface{} {
	if gs == nil {
		return nil
	}
	cp := &GPUDevices{
		Name:   gs.Name,
		Device: make(map[int]*GPUDevice, len(gs.Device)),
	}
	for id, dev := range gs.Device {
		newDev := &GPUDevice{
			ID:     dev.ID,
			Memory: dev.Memory,
			PodMap: make(map[string]*v1.Pod, len(dev.PodMap)),
		}
		for uid, pod := range dev.PodMap {
			newDev.PodMap[uid] = pod
		}
		cp.Device[id] = newDev
	}
	return cp
}

func (gs *GPUDevices) ScoreNode(pod *v1.Pod, schedulePolicy string) float64 {
	return 0
}

func (gs *GPUDevices) Allocate(kubeClient kubernetes.Interface, pod *v1.Pod) (*devices.DeviceReservation, error) {
	klog.V(4).Infoln("DeviceSharing:Into AllocateToPod", pod.Name)
	if getGPUMemoryOfPod(pod) > 0 {
		if NodeLockEnable {
			nodelock.UseClient(kubeClient)
			err := nodelock.LockNode(gs.Name, "gpu")
			if err != nil {
				return nil, errors.Errorf("node %s locked for lockname gpushare %s", gs.Name, err.Error())
			}
		}
		ids := predicateGPUbyMemory(pod, gs)
		if len(ids) == 0 {
			return nil, errors.Errorf("the node %s can't place the pod %s in ns %s", pod.Spec.NodeName, pod.Name, pod.Namespace)
		}
		id := ids[0]
		dev, ok := gs.Device[id]
		if !ok {
			return nil, errors.Errorf("failed to get GPU %d from node %s", id, gs.Name)
		}
		dev.PodMap[string(pod.UID)] = pod
		klog.V(4).Infof("predicates with gpu sharing, update pod %s/%s allocate to node [%s]", pod.Namespace, pod.Name, gs.Name)
		return &devices.DeviceReservation{
			DeviceType:  DeviceName,
			Annotations: gpuIndexAnnotations([]int{id}),
		}, nil
	}
	if getGPUNumberOfPod(pod) > 0 {
		ids := predicateGPUbyNumber(pod, gs)
		if len(ids) == 0 {
			return nil, errors.Errorf("the node %s can't place the pod %s in ns %s", pod.Spec.NodeName, pod.Name, pod.Namespace)
		}
		for _, id := range ids {
			dev, ok := gs.Device[id]
			if !ok {
				return nil, errors.Errorf("failed to get GPU %d from node %s", id, gs.Name)
			}
			dev.PodMap[string(pod.UID)] = pod
		}
		klog.V(4).Infof("predicates with gpu number, update pod %s/%s allocate to node [%s]", pod.Namespace, pod.Name, gs.Name)
		return &devices.DeviceReservation{
			DeviceType:  DeviceName,
			Annotations: gpuIndexAnnotations(ids),
		}, nil
	}
	return nil, nil
}
