package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"strconv"

	"google.golang.org/grpc"
	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
)

const (
	vcoreEndpoint        = "/var/lib/kubelet/device-plugins/vcore.sock"
	computeResourceName  = "nvidia-mps.com/vcore"
	computePartitionSize = 1
)

type vcoreManager struct {
	deviceManager
}

func NewVcoreServer(mm *mpsManager) DeviceServer {
	return vcoreManager{
		deviceManager: deviceManager{
			grpcServer: grpc.NewServer(),
			manager:    mm,
		},
	}
}

func (vm vcoreManager) DeviceName() string {
	log.Println("vcore-DeviceName() called ...")
	return computeResourceName
}

func (vm vcoreManager) DeviceEndpoint() string {
	log.Println("vcore-DeviceEndpoint() called ...")
	return vcoreEndpoint
}

func (vm vcoreManager) Run() error {
	log.Println("vcore-Run() called ...")
	log.Println("starting vcore-server ...")
	lis, err := net.Listen("unix", vcoreEndpoint)
	if err != nil {
		log.Println("starting vcore-server failed : %v", err)
	}
	vm.socket = vcoreEndpoint
	vm.grpcServer = grpc.NewServer()

	pluginapi.RegisterDevicePluginServer(vm.grpcServer, vm)
	return vm.grpcServer.Serve(lis)
}

func (vm vcoreManager) Stop() {
	log.Println("vcore-Stop() called ...")
	vm.grpcServer.Stop()
}

func (vm vcoreManager) ListDevices() []*pluginapi.Device {
	log.Println("vcore-ListDevices() called ...")
	computeDevicesCount := 100 / computePartitionSize
	virtualDevices := make([]*pluginapi.Device, computeDevicesCount)

	for i := 0; i < computeDevicesCount; i++ {
		virtualDeviceID := fmt.Sprintf("%s-%d", computeResourceName, i)
		virtualDevices[i] = &pluginapi.Device{
			ID:     virtualDeviceID,
			Health: pluginapi.Healthy,
		}
	}
	return virtualDevices
}

func (vm vcoreManager) ListAndWatch(empty *pluginapi.Empty, stream pluginapi.DevicePlugin_ListAndWatchServer) error {
	log.Println("vcore-ListAndWatch() called ...")
	resp := new(pluginapi.ListAndWatchResponse)
	resp.Devices = vm.ListDevices()
	vm.devices = resp.Devices
	if err := stream.Send(resp); err != nil {
		log.Println("error sending device list of vcores ... %v", err)
		return err
	}
	log.Println("successfully sent the list of devices of vcores ...")
	select {}
}

func (vm vcoreManager) Allocate(ctx context.Context, rqt *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	log.Println("vcore-Allocate() called ...")
	allocateResponse := &pluginapi.AllocateResponse{}

	for _, req := range rqt.ContainerRequests {
		containerAllocateResponse := &pluginapi.ContainerAllocateResponse{}
		totalCompute := len(req.DevicesIDs)
		envs := make(map[string]string)
		envs["CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"] = strconv.Itoa(totalCompute)
		envs["CUDA_MPS_PINNED_DEVICE_MEM_LIMIT"] = "0=2G"
		mounts := []*pluginapi.Mount{
			{
				ContainerPath: "/usr/local/nvidia",
				HostPath:      "/usr/local/nvidia",
			},
		}
		containerAllocateResponse.Envs = envs
		containerAllocateResponse.Mounts = mounts
		allocateResponse.ContainerResponses = append(allocateResponse.ContainerResponses, containerAllocateResponse)
		log.Println("successfully allocated the devices for vcore ...")
	}
	return allocateResponse, nil
}

func (vm vcoreManager) GetDevicePluginOptions(context.Context, *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
	log.Println("vcore-GetDevicePluginOptions() called ...")
	return &pluginapi.DevicePluginOptions{
		PreStartRequired:                false,
		GetPreferredAllocationAvailable: false,
	}, nil
}

func (vm vcoreManager) GetPreferredAllocation(ctx context.Context, rqt *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
	log.Println("vcore-getPreferredAllocation() called ...")
	preferredAllocateResponse := &pluginapi.PreferredAllocationResponse{}
	return preferredAllocateResponse, nil
}

func (vm vcoreManager) PreStartContainer(context.Context, *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	log.Println("vcore-PreStartContainer() called ...")
	preStartContainerResponse := pluginapi.PreStartContainerResponse{}
	return &preStartContainerResponse, nil
}
