package plugin

import (
	"bufio"
	"fmt"
	"github.com/golang/glog"
	"github.com/kubevirt/device-plugin-manager/pkg/dpm"
	"golang.org/x/net/context"
	"k8s-device-plugin/internal/pkg/allocator"
	"k8s-device-plugin/internal/pkg/hydcu"
	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
)

type DCULister struct {
	ResUpdateChan chan dpm.PluginNameList
	Heartbeat     chan bool
	Signal        chan os.Signal
}

type DCUPlugin struct {
	DCUs               map[string]map[string]interface{}
	Heartbeat          chan bool
	signal             chan os.Signal
	Resource           string
	devAllocator       allocator.Policy
	allocatorInitError bool
}
type DCUPluginOption func(*DCUPlugin)

func (p *DCUPlugin) Start() error {
	p.signal = make(chan os.Signal, 1)
	signal.Notify(p.signal, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)
	err := p.devAllocator.Init(getDevices(), "")
	if err != nil {
		glog.Errorf("allocator init failed. Falling back to kubelet default allocation. Error %v", err)
		p.allocatorInitError = true
	}
	return nil
}
func getDevices() []*allocator.Device {
	devices := hydcu.GetHYDCUs()
	var deviceList []*allocator.Device

	for id, deviceData := range devices {
		device := &allocator.Device{
			Id:                   id,
			Card:                 deviceData["card"].(int),
			RenderD:              deviceData["renderD"].(int),
			DevId:                deviceData["devID"].(string),
			ComputePartitionType: deviceData["computePartitionType"].(string),
			MemoryPartitionType:  deviceData["memoryPartitionType"].(string),
			NodeId:               deviceData["nodeId"].(int),
			NumaNode:             deviceData["numaNode"].(int),
		}
		deviceList = append(deviceList, device)
	}
	return deviceList
}

// Stop is an optional interface that could be implemented by plugin.
// If case Stop is implemented, it will be executed by Manager after the
// plugin is unregistered from kubelet. This method could be used to tear
// down resources.
func (p *DCUPlugin) Stop() error {
	return nil
}

func NewDCUPlugin(options ...DCUPluginOption) *DCUPlugin {
	dcuPlugin := &DCUPlugin{}
	for _, option := range options {
		option(dcuPlugin)
	}
	return dcuPlugin
}

func WithAllocator(a allocator.Policy) DCUPluginOption {
	return func(p *DCUPlugin) {
		p.devAllocator = a
	}
}

func WithHeartbeat(ch chan bool) DCUPluginOption {
	return func(p *DCUPlugin) {
		p.Heartbeat = ch
	}
}

func WithResource(res string) DCUPluginOption {
	return func(p *DCUPlugin) {
		p.Resource = res
	}
}
func (l *DCULister) GetResourceNamespace() string {
	return "c-3000.com"
}

func (l *DCULister) Discover(pluginListCh chan dpm.PluginNameList) {
	for {
		select {
		case newResourcesList := <-l.ResUpdateChan:
			pluginListCh <- newResourcesList
		case <-pluginListCh: // Stop message received
			// Stop resourceUpdateCh
			return
		}
	}
}

func (p *DCUPlugin) GetDevicePluginOptions(ctx context.Context, e *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
	if p.allocatorInitError {
		return &pluginapi.DevicePluginOptions{}, nil
	}
	return &pluginapi.DevicePluginOptions{
		GetPreferredAllocationAvailable: true,
	}, nil
}

func (p *DCUPlugin) PreStartContainer(ctx context.Context, r *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	return &pluginapi.PreStartContainerResponse{}, nil
}

func simpleHealthCheck() bool {
	entries, err := filepath.Glob("/sys/class/kfd/kfd/topology/nodes/*/properties")
	if err != nil {
		glog.Errorf("Error finding properties files: %v", err)
		return false
	}

	for _, propFile := range entries {
		f, err := os.Open(propFile)
		if err != nil {
			glog.Errorf("Error opening %s: %v", propFile, err)
			continue
		}

		defer f.Close()

		var cpuCores, gfxVersion int
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, "cpu_cores_count") {
				parts := strings.Fields(line)
				if len(parts) == 2 {
					cpuCores, _ = strconv.Atoi(parts[1])
				}
			} else if strings.HasPrefix(line, "gfx_target_version") {
				parts := strings.Fields(line)
				if len(parts) == 2 {
					gfxVersion, _ = strconv.Atoi(parts[1])
				}
			}
		}

		if err := scanner.Err(); err != nil {
			glog.Warningf("Error scanning %s: %v", propFile, err)
			continue
		}

		if cpuCores == 0 && gfxVersion > 0 {
			return true
		}
	}
	glog.Warning("No GPU nodes found via properties")
	return false
}

func (p *DCUPlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {

	p.DCUs = hydcu.GetHYDCUs()

	glog.Infof("Found %d DCUs", len(p.DCUs))

	devs := make([]*pluginapi.Device, len(p.DCUs))
	var isHomogeneous bool
	isHomogeneous = hydcu.IsHomogeneous()
	// Initialize a map to store partitionType based device list
	resourceTypeDevs := make(map[string][]*pluginapi.Device)

	if isHomogeneous {
		// limit scope for hwloc
		func() {
			i := 0
			for id, device := range p.DCUs {
				dev := &pluginapi.Device{
					ID:     id,
					Health: pluginapi.Healthy,
				}
				devs[i] = dev
				i++

				numas := []int64{int64(device["numaNode"].(int))}
				glog.Infof("Watching GPU with bus ID: %s NUMA Node: %+v", id, numas)

				numaNodes := make([]*pluginapi.NUMANode, len(numas))
				for j, v := range numas {
					numaNodes[j] = &pluginapi.NUMANode{
						ID: int64(v),
					}
				}

				dev.Topology = &pluginapi.TopologyInfo{
					Nodes: numaNodes,
				}
			}
		}()
		s.Send(&pluginapi.ListAndWatchResponse{Devices: devs})
	} else {
		func() {
			for id, device := range p.DCUs {
				dev := &pluginapi.Device{
					ID:     id,
					Health: pluginapi.Healthy,
				}
				// Append a device belonging to a certain partition type to its respective list
				partitionType := device["computePartitionType"].(string) + "_" + device["memoryPartitionType"].(string)
				resourceTypeDevs[partitionType] = append(resourceTypeDevs[partitionType], dev)

				numas := []int64{int64(device["numaNode"].(int))}
				glog.Infof("Watching GPU with bus ID: %s NUMA Node: %+v", id, numas)

				numaNodes := make([]*pluginapi.NUMANode, len(numas))
				for j, v := range numas {
					numaNodes[j] = &pluginapi.NUMANode{
						ID: int64(v),
					}
				}

				dev.Topology = &pluginapi.TopologyInfo{
					Nodes: numaNodes,
				}
			}
		}()
		// Send the appropriate list of devices based on the partitionType
		if devList, exists := resourceTypeDevs[p.Resource]; exists {
			s.Send(&pluginapi.ListAndWatchResponse{Devices: devList})
		}
	}

loop:
	for {
		select {
		case <-p.Heartbeat:
			var health = pluginapi.Unhealthy

			if simpleHealthCheck() {
				health = pluginapi.Healthy
			}

			// update with per device GPU health status
			if isHomogeneous {
				//exporter.PopulatePerGPUDHealth(devs, health)
				glog.Infof("Healthy: %v", health)
				s.Send(&pluginapi.ListAndWatchResponse{Devices: devs})
			} else {
				if devList, exists := resourceTypeDevs[p.Resource]; exists {
					//exporter.PopulatePerGPUDHealth(devList, health)
					s.Send(&pluginapi.ListAndWatchResponse{Devices: devList})
				}
			}

		case <-p.signal:
			glog.Infof("Received signal, exiting")
			break loop
		}
	}
	// returning a value with this function will unregister the plugin from k8s

	return nil
}

func (p *DCUPlugin) GetPreferredAllocation(ctx context.Context, req *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
	response := &pluginapi.PreferredAllocationResponse{}
	for _, req := range req.ContainerRequests {
		allocated_ids, err := p.devAllocator.Allocate(req.AvailableDeviceIDs, req.MustIncludeDeviceIDs, int(req.AllocationSize))
		if err != nil {
			glog.Errorf("unable to get preferred allocation list. Error:%v", err)
			return nil, fmt.Errorf("unable to get preferred allocation list. Error:%v", err)
		}
		resp := &pluginapi.ContainerPreferredAllocationResponse{
			DeviceIDs: allocated_ids,
		}
		response.ContainerResponses = append(response.ContainerResponses, resp)
	}
	return response, nil
}

func (p *DCUPlugin) Allocate(ctx context.Context, r *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	var response pluginapi.AllocateResponse
	var car pluginapi.ContainerAllocateResponse
	var dev *pluginapi.DeviceSpec

	for _, req := range r.ContainerRequests {
		car = pluginapi.ContainerAllocateResponse{}

		// Currently, there are only 1 /dev/kfd per nodes regardless of the # of GPU available
		// for compute/rocm/HSA use cases
		dev = new(pluginapi.DeviceSpec)
		dev.HostPath = "/dev/kfd"
		dev.ContainerPath = "/dev/kfd"
		dev.Permissions = "rw"
		car.Devices = append(car.Devices, dev)

		for _, id := range req.DevicesIDs {
			glog.Infof("Allocating device ID: %s", id)

			for k, v := range p.DCUs[id] {
				// Map struct previously only had 'card' and 'renderD' and only those are paths to be appended as before
				if k != "card" && k != "renderD" {
					continue
				}
				devpath := fmt.Sprintf("/dev/dri/%s%d", k, v)
				dev = new(pluginapi.DeviceSpec)
				dev.HostPath = devpath
				dev.ContainerPath = devpath
				dev.Permissions = "rw"
				car.Devices = append(car.Devices, dev)
			}
		}

		response.ContainerResponses = append(response.ContainerResponses, &car)
	}

	return &response, nil
}

func (l *DCULister) NewPlugin(resourceLastName string) dpm.PluginInterface {
	options := []DCUPluginOption{
		WithHeartbeat(l.Heartbeat),
		WithResource(resourceLastName),
		WithAllocator(allocator.NewBestEffortPolicy()),
	}
	return NewDCUPlugin(options...)
}
