package main

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

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

const (
	kubeletEndpoint = "/var/lib/kubelet/device-plugins/kubelet.sock"
)

type deviceManager struct {
	grpcServer *grpc.Server
	devices    []*pluginapi.Device
	socket     string
	manager    *mpsManager
}

type DeviceServer interface {
	Run() error
	Stop()
	DeviceEndpoint() string
	DeviceName() string
}

type mpsManager struct {
	deviceServers map[string]DeviceServer
	grpcServer    *grpc.Server
}

func NewMPSManager() mpsManager {
	return mpsManager{
		deviceServers: make(map[string]DeviceServer),
		grpcServer:    grpc.NewServer(),
	}
}

func (mm *mpsManager) RunGRPCService() {
	vcoreServer := NewVcoreServer(mm)
	mm.deviceServers[vcoreServer.DeviceName()] = vcoreServer
	for name, srv := range mm.deviceServers {
		log.Printf("server %s is running", name)
		go srv.Run()
	}
}

func (mm *mpsManager) RegisterToKubelet() error {
	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)

	for _, srv := range mm.deviceServers {
		request := &pluginapi.RegisterRequest{
			Version:      pluginapi.Version,
			Endpoint:     srv.DeviceEndpoint(),
			ResourceName: srv.DeviceName(),
		}
		if _, err = client.Register(context.Background(), request); err != nil {
			return fmt.Errorf("device plugin cannot register to kubelet service: %v", err)
		}
	}
	return nil
}
