package main

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

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

type GPUDevicePlugin struct {
	socket string
	server *grpc.Server
}

func newGPUDevicePlugin(socket string) *GPUDevicePlugin {
	return &GPUDevicePlugin{
		socket: socket,
	}
}

func (devicePlugin *GPUDevicePlugin) Start() error {
	lis, err := net.Listen("unix", devicePlugin.socket)
	if err != nil {
		return err
	}
	devicePlugin.server = grpc.NewServer()
	pluginapi.RegisterDevicePluginServer(devicePlugin.server, devicePlugin)
	log.Println("Starting GPU Device plugin server ...")
	return devicePlugin.server.Serve(lis)
}

func (devicePlugin *GPUDevicePlugin) Stop() {
	if devicePlugin.server != nil {
		devicePlugin.server.Stop()
	}
}

func (devicePlugin *GPUDevicePlugin) ListAndWatch(empty *pluginapi.Empty, stream pluginapi.DevicePlugin_ListAndWatchServer) error {
	resp := new(pluginapi.ListAndWatchResponse)
	devices := make([]*pluginapi.Device, 104)
	for i := 0; i < 100; i++ {
		devices[i] = &pluginapi.Device{
			ID:     "mps.iitb/gpu-vcore-" + strconv.Itoa(i),
			Health: pluginapi.Healthy,
		}
	}
	for i := 0; i < 4; i++ {
		devices[i] = &pluginapi.Device{
			ID:     "mps.iitb/gpu-vmemory-" + strconv.Itoa(i),
			Health: pluginapi.Healthy,
		}
	}
	resp.Devices = devices
	if err := stream.Send(resp); err != nil {
		log.Printf("Error sending device list : %v", err)
		return err
	}
	log.Printf("Successfully sent the list of devices ...")
	select {}
}

func (devicePlugin *GPUDevicePlugin) Allocate(ctx context.Context, rqt *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	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.Printf("Successfully allocated the devices ...")
	}
	return allocateResponse, nil
}

func (devicePlugin *GPUDevicePlugin) GetDevicePluginOptions(context.Context, *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
	return &pluginapi.DevicePluginOptions{
		PreStartRequired:                false,
		GetPreferredAllocationAvailable: false,
	}, nil
}

func (devicePlugin *GPUDevicePlugin) GetPreferredAllocation(ctx context.Context, rqt *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
	preferredAllocateResponse := &pluginapi.PreferredAllocationResponse{}

	for _, req := range rqt.ContainerRequests {
		containerAllocateResponse := &pluginapi.ContainerPreferredAllocationResponse{}
		totalDevicesToAllot := req.AllocationSize
		for i := 0; i < len(req.MustIncludeDeviceIDs); i++ {
			containerAllocateResponse.DeviceIDs = append(containerAllocateResponse.DeviceIDs, req.MustIncludeDeviceIDs[i])
			totalDevicesToAllot--
		}
		preferredAllocateResponse.ContainerResponses = append(preferredAllocateResponse.ContainerResponses, containerAllocateResponse)
	}
	return preferredAllocateResponse, nil
}

func (devicePlugin *GPUDevicePlugin) PreStartContainer(context.Context, *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	preStartContainerResponse := pluginapi.PreStartContainerResponse{}
	return &preStartContainerResponse, nil
}

func main() {
	socketPlugin := "/var/lib/kubelet/device-plugins/gpu-device-plugin.sock"
	devicePlugin := newGPUDevicePlugin(socketPlugin)
	defer devicePlugin.Stop()

	if err := devicePlugin.Start(); err != nil {
		log.Fatalf("Error starting the device plugin : %v", err)
	}
}
