package wait

import (
	"context"
	"os"
	"regexp"
	"time"

	"edge-infra.dev/pkg/lib/fog"
	v1deviceclass "edge-infra.dev/pkg/sds/devices/k8s/apis/v1"
)

const (
	deviceDir       = "/var/lib/kubelet/device-plugins"
	deviceCachePath = "/zynstra/config"
)

var (
	checkInterval = time.Second * 1
	pattern       = regexp.MustCompile(`^ds-.*\.sock$`)
)

// ForDevices waits until a desired number of devices are registered and the
// registration process has stabilized (devices are no longer being registered)
func ForDevices(ctx context.Context) error {
	w := waiter{
		deviceDir:       deviceDir,
		deviceCachePath: deviceCachePath,
		checkInterval:   checkInterval,
	}

	return w.wait(ctx)
}

// waiter contains of the configuration options for the process of waiting for
// device registration to complete
type waiter struct {
	deviceDir       string
	deviceCachePath string
	checkInterval   time.Duration
}

// wait for the device registration process to complete. This method will not
// return until device registration has successfully completed
func (w *waiter) wait(ctx context.Context) error {
	log := fog.FromContext(ctx)

	log.Info("waiting for devices...", "deviceDirectory", w.deviceDir, "deviceCacheDir", w.deviceCachePath, "deviceCacheFilename", "deviceclasses.json")

	deviceClasses, err := v1deviceclass.ReadDeviceClassFromPersistence(w.deviceCachePath)
	if err != nil {
		return err
	}

	expectedDeviceCount := len(deviceClasses)
	for {
		registrationComplete, err := w.checkDevices(ctx, expectedDeviceCount)
		if err != nil {
			return err
		}

		if registrationComplete {
			log.Info("device registration complete")
			break
		}
		time.Sleep(w.checkInterval)
	}
	return nil
}

// checkDevices returns true if the expected number of devices have been
// registered
func (w *waiter) checkDevices(ctx context.Context, expectedDeviceCount int) (bool, error) {
	log := fog.FromContext(ctx)

	files, err := os.ReadDir(w.deviceDir)
	if err != nil {
		return false, err
	}

	count := 0
	for _, file := range files {
		if !file.IsDir() && pattern.MatchString(file.Name()) {
			count++
		}
	}
	log.Info("found devices", "expectedCount", expectedDeviceCount, "count", count)
	return count >= expectedDeviceCount, nil
}