package main

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

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

const (
	mpsControlBin      = "/usr/local/nvidia/bin/nvidia-cuda-mps-control"
	mpsActiveThreadCmd = "get_default_active_thread_percentage"
	mpsMemLimitEnv     = "CUDA_MPS_PINNED_DEVICE_MEM_LIMIT"
	mpsThreadLimitEnv  = "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"
	socketPlugin       = "/var/lib/kubelet/device-plugins/mps-device-plugin.sock"
)

var (
	computeResourceName = "nvidia-mps.com/vcore"
	memResourceName     = "nvidia-mps.com/vmemory"
)

type mpsGPUManager struct {
	grpcServer           *grpc.Server
	devices              map[string]pluginapi.Device
	socket               string
	computePartitionSize int // In utilization (eg. 1%)
	memPartitionSize     int // In MB (eg. 256MB)
}

func NewMpsGPUManager(computePartitionSize, memPartitionSize int) *mpsGPUManager {
	return &mpsGPUManager{
		devices:              make(map[string]pluginapi.Device),
		computePartitionSize: computePartitionSize,
		memPartitionSize:     memPartitionSize,
		socket:               socketPlugin,
	}
}

func (mgm *mpsGPUManager) ListDevices() []*pluginapi.Device {
	gpuMemoryAvailable := 16384 // Using static value for now
	computeDevicesCount := 100 / mgm.computePartitionSize
	memoryDevicesCount := gpuMemoryAvailable / mgm.memPartitionSize
	virtualDevices := make([]*pluginapi.Device, computeDevicesCount+memoryDevicesCount)

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

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

// func (mgm *mpsGPUManager) isMpsHealthy() error {
// 	glog.Infof("Checking status of MPS on the node")

// 	var out bytes.Buffer
// 	reader, writer := io.Pipe()
// 	defer writer.Close()
// 	defer reader.Close()

// 	mpsCmd := exec.Command(mpsControlBin)
// 	mpsCmd.Stdin = reader
// 	mpsCmd.Stdout = &out

// 	err := mpsCmd.Start()
// 	if err != nil {
// 		return fmt.Errorf("failed to start NVIDIA MPS health check: %v", err)
// 	}

// 	writer.Write([]byte(mpsActiveThreadCmd))
// 	writer.Close()

// 	err = mpsCmd.Wait()
// 	if err != nil {
// 		return fmt.Errorf("failed to health check NVIDIA MPS: %v", err)
// 	}

// 	reader.Close()
// 	glog.Infof("MPS is healthy, active thread percentage = %s", out.String())
// 	return nil
// }

func (mgm *mpsGPUManager) Start() error {
	glog.Infof("Starting MPS GPU Manager")
	lis, err := net.Listen("unix", mgm.socket)
	if err != nil {
		return err
	}
	mgm.grpcServer = grpc.NewServer()
	pluginapi.RegisterDevicePluginServer(mgm.grpcServer, mgm)
	glog.Infof("MPS GPU Manager registered with the kubelet")
	return mgm.grpcServer.Serve(lis)
}

func (mgm *mpsGPUManager) Stop() {
	if mgm.grpcServer != nil {
		mgm.grpcServer.Stop()
	}
	glog.Infof("MPS GPU Manager stopped")
}

func (mgm *mpsGPUManager) ListAndWatch(empty *pluginapi.Empty, stream pluginapi.DevicePlugin_ListAndWatchServer) error {
	resp := new(pluginapi.ListAndWatchResponse)
	resp.Devices = mgm.ListDevices()
	if err := stream.Send(resp); err != nil {
		glog.Infof("Error sending device list : %v", err)
		return err
	}
	glog.Infof("Successfully sent the list of devices ...")
	select {}
}

func (mgm *mpsGPUManager) 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)
		glog.Infof("Successfully allocated the devices ...")
	}
	return allocateResponse, nil
}

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

func (mgm *mpsGPUManager) 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 (mgm *mpsGPUManager) PreStartContainer(context.Context, *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	preStartContainerResponse := pluginapi.PreStartContainerResponse{}
	return &preStartContainerResponse, nil
}

func main() {
	mgm := NewMpsGPUManager(1, 256)
	defer mgm.Stop()

	if err := mgm.Start(); err != nil {
		glog.Fatalf("Error starting the MPS GPU Manager : %v", err)
	}
}
