package main

import (
	"flag"
	"fmt"
	"github.com/golang/glog"
	"github.com/kubevirt/device-plugin-manager/pkg/dpm"
	"k8s-device-plugin/internal/pkg/hwloc"
	"k8s-device-plugin/internal/pkg/hydcu"
	"k8s-device-plugin/internal/pkg/plugin"
	"os"
	"time"
)

var gitDescribe string

type ResourceNamingStrategy string

const (
	StrategySingle ResourceNamingStrategy = "single"
	StrategyMixed  ResourceNamingStrategy = "mixed"
)

func ParseStrategy(s string) (ResourceNamingStrategy, error) {
	switch s {
	case string(StrategySingle):
		return StrategySingle, nil
	case string(StrategyMixed):
		return StrategyMixed, nil
	default:
		return "", fmt.Errorf("invalid strategy: %s", s)
	}
}

// return []string{"dcu"}
func getResourceList(resourceNamingStrategy ResourceNamingStrategy) ([]string, error) {
	var resources []string

	// Check if the node is homegeneous
	isHomogeneous := hydcu.IsHomogeneous()
	partitionCountMap := hydcu.UniquePartitionConfigCount(hydcu.GetHYDCUs())
	if len(hydcu.GetHYDCUs()) == 0 {
		return resources, nil
	}
	if isHomogeneous {
		if resourceNamingStrategy == StrategySingle {
			resources = []string{"dcu"}
		} else if resourceNamingStrategy == StrategyMixed {
			if len(partitionCountMap) == 0 {
				// If partitioning is not supported on the node, we should report resources under "gpu" regardless of the strategy
				resources = []string{"dcu"}
			} else {
				for partitionType, count := range partitionCountMap {
					if count > 0 {
						resources = append(resources, partitionType)
					}
				}
			}
		}
	} else {
		if resourceNamingStrategy == StrategySingle {
			return resources, fmt.Errorf("Partitions of different styles across DCUs in a node is not supported with single strategy. Please start device plugin with mixed strategy")
		} else if resourceNamingStrategy == StrategyMixed {
			for partitionType, count := range partitionCountMap {
				if count > 0 {
					resources = append(resources, partitionType)
				}
			}
		}
	}
	return resources, nil
}
func main() {
	versions := [...]string{
		"DCU device plugin for Kubernetes",
		fmt.Sprintf("%s version %s", os.Args[0], gitDescribe),
		fmt.Sprintf("%s", hwloc.GetVersions()),
	}

	flag.Usage = func() {
		for _, v := range versions {
			fmt.Println(v)
		}
		fmt.Fprintln(os.Stderr, "Usage:")
		flag.PrintDefaults()
	}
	var pulse int
	var resourceNamingStrategy string
	flag.IntVar(&pulse, "pulse", 0, "time between health check polling in seconds. Set to 0 to disable health check")
	flag.StringVar(&resourceNamingStrategy, "resource-naming-strategy", "single", "Resource strategy to be used: single or mixed")
	flag.Parse()
	strategy, err := ParseStrategy(resourceNamingStrategy)
	if err != nil {
		glog.Errorf("%v", err)
	}

	for _, v := range versions {
		glog.Infof("%s", v)
	}

	l := plugin.DCULister{
		ResUpdateChan: make(chan dpm.PluginNameList),
		Heartbeat:     make(chan bool),
	}

	manager := dpm.NewManager(&l)

	if pulse > 0 {
		go func() {
			glog.Infof("Heart beating every %d seconds", pulse)
			for {
				time.Sleep(time.Second * time.Duration(pulse))
				l.Heartbeat <- true
			}
		}()
	}

	go func() {
		// /sys/class/kfd only exists if ROCm kernel/driver is installed
		var path = "/sys/class/kfd"
		if _, err := os.Stat(path); err == nil {
			resources, err := getResourceList(strategy)
			if err != nil {
				glog.Errorf("Error occured: %v", err)
				os.Exit(1)
			}
			if len(resources) > 0 {
				l.ResUpdateChan <- resources
			}
		}
	}()
	manager.Run()
}
