package main

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

	"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"
	pluginEndpoint     = "/device-plugin/mps-device-plugin.sock"
	kubeletEndpoint    = "/device-plugin/kubelet.sock"
)

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

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

func NewMpsGPUManager(computePartitionSize int) *mpsGPUManager {
	return &mpsGPUManager{
		devices:              make(map[string]pluginapi.Device),
		computePartitionSize: computePartitionSize,
		stop:                 make(chan bool),
	}
}

type pluginService struct {
	mgm *mpsGPUManager
}

// 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) Serve() {
	log.Println("Serve() called ...")
	lis, err := net.Listen("unix", pluginEndpoint)
	if err != nil {
		log.Fatalf("starting device plugin server failed : %v", err)
	}
	mgm.socket = pluginEndpoint
	mgm.grpcServer = grpc.NewServer()

	pluginbeta := &pluginService{mgm: mgm}
	pluginbeta.RegisterService()
	registeredWithKubelet := false

	for {
		select {
		case <-mgm.stop:
			close(mgm.stop)
			return
		default:
			{
				var wg sync.WaitGroup
				wg.Add(1)
				// Starts device plugin service.
				go func() {
					defer wg.Done()
					// Blocking call to accept incoming connections.
					err := mgm.grpcServer.Serve(lis)
					log.Fatalf("device-plugin server stopped serving: %v", err)
				}()

				if !registeredWithKubelet {
					for len(mgm.grpcServer.GetServiceInfo()) <= 0 {
						time.Sleep(1 * time.Second)
					}
					log.Println("device-plugin server started serving")

					err = RegisterToKubelet()
					if err != nil {
						mgm.grpcServer.Stop()
						wg.Wait()
						glog.Fatal(err)
					}
					log.Println("device plugin registered with kubelet")
					registeredWithKubelet = true
				}
			}
		}
	}
}

func RegisterToKubelet() error {
	log.Println("RegisterToKubelet() called ...")
	conn, err := grpc.Dial(kubeletEndpoint, grpc.WithInsecure(),
		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
			return net.DialTimeout("unix", addr, timeout)
		}))
	if err != nil {
		return fmt.Errorf("device plugin cannot connect to the kubelet service: %v", err)
	}
	defer conn.Close()
	client := pluginapi.NewRegistrationClient(conn)

	request := &pluginapi.RegisterRequest{
		Version:      pluginapi.Version,
		Endpoint:     "mps-device-plugin.sock",
		ResourceName: "nvidia.com/gpu",
	}

	if _, err = client.Register(context.Background(), request); err != nil {
		return fmt.Errorf("device plugin cannot register to kubelet service: %v", err)
	}
	return nil
}

func (ps *pluginService) RegisterService() {
	log.Println("RegisterSerivce() called ...")
	pluginapi.RegisterDevicePluginServer(ps.mgm.grpcServer, ps)
}

func (mgm *mpsGPUManager) Stop() {
	log.Println("Stop() called ...")
	if mgm.grpcServer != nil {
		mgm.grpcServer.Stop()
	}
	mgm.stop <- true
	<-mgm.stop
	log.Println("MPS GPU Manager stopped")
}

func (ps *pluginService) ListDevices() []*pluginapi.Device {
	log.Println("ListDevices() called ...")
	computeDevicesCount := 100 / ps.mgm.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 (ps *pluginService) ListAndWatch(empty *pluginapi.Empty, stream pluginapi.DevicePlugin_ListAndWatchServer) error {
	log.Println("ListAndWatch() called ...")
	resp := new(pluginapi.ListAndWatchResponse)
	resp.Devices = ps.ListDevices()
	if err := stream.Send(resp); err != nil {
		log.Fatalf("error sending device list : %v", err)
		return err
	}
	log.Println("successfully sent the list of devices ...")
	select {}
}

func (ps *pluginService) Allocate(ctx context.Context, rqt *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	log.Println("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_PIPE_DIRECTORY"] = "/tmp/nvidia-mps"
		envs["CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"] = strconv.Itoa(totalCompute)
		envs["CUDA_MPS_PINNED_DEVICE_MEM_LIMIT"] = "0=2G"
		envs["NVIDIA_DRIVER_CAPABILITIES"] = "compute,utility"
		envs["NVIDIA_RUNTIME"] = "nvidia"
		envs["NVIDIA_VERSION"] = "11.7"
		envs["NVIDIA_VISIBLE_DEVICES"] = "0"

		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidiactl",
			ContainerPath: "/dev/nvidiactl",
			Permissions:   "mrw",
		})
		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidia-modeset",
			ContainerPath: "/dev/nvidia-modeset",
			Permissions:   "mrw",
		})
		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidia-uvm",
			ContainerPath: "/dev/nvidia-uvm",
			Permissions:   "mrw",
		})
		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidia-uvm-tools",
			ContainerPath: "/dev/nvidia-uvm-tools",
			Permissions:   "mrw",
		})
		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidia0",
			ContainerPath: "/dev/nvidia0",
			Permissions:   "mrw",
		})
		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidia-caps/nvidia-cap1",
			ContainerPath: "/dev/nvidia-caps/nvidia-cap1",
			Permissions:   "mrw",
		})
		containerAllocateResponse.Devices = append(containerAllocateResponse.Devices, &pluginapi.DeviceSpec{
			HostPath:      "/dev/nvidia-caps/nvidia-cap2",
			ContainerPath: "/dev/nvidia-caps/nvidia-cap2",
			Permissions:   "mrw",
		})
		mounts := []*pluginapi.Mount{
			{
				ContainerPath: "/tmp/nvidia-mps",
				HostPath:      "/tmp/nvidia-mps",
				ReadOnly:      false,
			},
			{
				ContainerPath: "/usr/local/nvidia",
				HostPath:      "/usr/src/nvidia-515.76",
				ReadOnly:      false,
			},
		}
		log.Println("Mounted the required directories")
		containerAllocateResponse.Envs = envs
		containerAllocateResponse.Mounts = mounts
		allocateResponse.ContainerResponses = append(allocateResponse.ContainerResponses, containerAllocateResponse)
		log.Println("successfully allocated the required devices ...")
	}
	return allocateResponse, nil
}

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

func (ps *pluginService) GetPreferredAllocation(ctx context.Context, rqt *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
	log.Println("GetPreferredAllocation() called ...")
	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 (ps *pluginService) PreStartContainer(context.Context, *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	log.Println("PreStartContainer() called ...")
	preStartContainerResponse := pluginapi.PreStartContainerResponse{}
	return &preStartContainerResponse, nil
}

func main() {
	mgm := NewMpsGPUManager(1)
	defer mgm.Stop()
	mgm.Serve()
}
