//go:build linux

package events

import (
	"context"
	"errors"
	"regexp"
	"strings"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/containers"
	typeurl "github.com/containerd/typeurl/v2"
	"github.com/opencontainers/runtime-spec/specs-go"

	"edge-infra.dev/pkg/lib/kernel/udev"
	cc "edge-infra.dev/pkg/sds/devices/agent/common"
	devctrs "edge-infra.dev/pkg/sds/devices/agent/containers"
	"edge-infra.dev/pkg/sds/devices/agent/udevproxy"
	"edge-infra.dev/pkg/sds/devices/class"
	dsv1 "edge-infra.dev/pkg/sds/devices/k8s/apis/v1"
	"edge-infra.dev/pkg/sds/devices/logger"
)

var udevProxyFn = udevproxy.ReplayUEventsToContainers

const (
	defaultExecutablePath = "/usr/local/bin/device-system-executable"
	executablePathEnvVar  = "DEVICE_SYSTEM_EXECUTABLE_PATH"
)

var executablePathRegex, _ = regexp.Compile("^DEVICE_SYSTEM_EXECUTABLE_PATH=.*$")

// UDevEventConstructor returns a constructor function that instantiates a new device udev event.
// The constructor function will parse a uevent objcet and attempt to add/remove the device from device classes.
// If the uevent is remove, the udev event will be replayed immediately to containers requesting the device.
func UDevEventConstructor(ctrClient *containerd.Client, deviceClasses map[string]*dsv1.DeviceClass, allContainers map[string]*containers.Container) func(ctx context.Context, udevEvent *udev.UEvent) (DeviceEvent, error) {
	return func(ctx context.Context, udevEvent *udev.UEvent) (DeviceEvent, error) {
		if udevEvent == nil {
			return nil, errors.New("error, attempted to parse nil uevent")
		}
		switch udevEvent.Action {
		case udev.AddAction:
			addDeviceFromUEvent(ctx, udevEvent, deviceClasses)
			return newUDevEvent(ctx, ctrClient, udevEvent, allContainers, deviceClasses)
		case udev.RemoveAction:
			event, err := newUDevEvent(ctx, ctrClient, udevEvent, allContainers, deviceClasses)
			if err != nil {
				return nil, err
			}
			removeDeviceFromUEvent(udevEvent, deviceClasses)
			udevproxy.ReplayRemoveUEvents(ctx, []*udev.UEvent{udevEvent}, allContainers)
			return event, nil
		default:
			return newUDevEvent(ctx, ctrClient, udevEvent, allContainers, deviceClasses)
		}
	}
}

// newUDevEvent returns a new device event built from udev event
func newUDevEvent(ctx context.Context, ctrClient *containerd.Client, uevent *udev.UEvent, allContainers map[string]*containers.Container, allDeviceClasses map[string]*dsv1.DeviceClass) (DeviceEvent, error) {
	if uevent == nil {
		return nil, errors.New("uevent cannot be nil")
	}

	event := &udevEvent{
		event: &event{
			containers: map[string]*containers.Container{},
			postHookFn: func(context.Context) {},
			timestamp:  time.Now(),
		},
	}
	log := logger.FromContext(ctx)

	containerExecFns := map[string]func(){}
	for _, ctr := range allContainers {
		ctrCtx := devctrs.WithContainerLogger(ctx, ctr)
		if !ueventMatchesContainer(ctrCtx, uevent, ctr, allDeviceClasses) {
			continue
		}

		executablePath := fetchExecutablePath(ctrCtx, ctr)
		rootPath, err := devctrs.FetchContainerRootPath(ctrCtx, ctrClient, ctr)
		if err != nil {
			log.Debug("could not find container root path", "error", err)
		} else {
			containerExecFns[ctr.ID] = func() {
				devctrs.NewExecFn(ctrCtx, ctr.Labels[cc.AnnContainerName], rootPath, executablePath, uevent.EnvVars)
			}
		}

		ctr.Labels[class.DefaultClass] = requested
		event.event.containers[ctr.ID] = ctr
	}

	event.event.postHookFn = func(ctx context.Context) {
		udevProxyFn(ctx, []*udev.UEvent{uevent}, event.event.containers)
		for _, fn := range containerExecFns {
			fn()
		}
	}
	return event, nil
}

// fetchExecutablePath will attempt to parse DEVICE_SCRIPT_PATH environment variable
// from the container process call if present.
func fetchExecutablePath(ctx context.Context, ctr *containers.Container) string {
	log := logger.FromContext(ctx)

	spec := &specs.Spec{}
	if err := typeurl.UnmarshalTo(ctr.Spec, spec); err != nil {
		log.Debug("could not parse container spec", "container", ctr.Labels[cc.AnnContainerName], "error", err)
		return defaultExecutablePath
	}

	for _, envVar := range spec.Process.Env {
		if executablePathRegex.Match([]byte(envVar)) {
			splitPath := strings.Split(envVar, "=")
			if len(splitPath) != 2 {
				return defaultExecutablePath
			}
			return splitPath[1]
		}
	}
	return defaultExecutablePath
}

// addDeviceFromUEvent will take a uevent and attempt to match it against the devices classes. If it matches, the device will be added to the device class.
func addDeviceFromUEvent(ctx context.Context, udevEvent *udev.UEvent, deviceClasses map[string]*dsv1.DeviceClass) {
	log := logger.FromContext(ctx)
	devicePath := prefixSysPath(udevEvent.SysPath)
	for _, devClass := range deviceClasses {
		if _, err := devClass.AddDeviceIfMatched(devicePath); err != nil {
			log.Error("error adding device to class", "class", devClass.ClassName(), "path", udevEvent.SysPath, "error", err)
			continue
		}
	}
}

// removeDeviceFromUEvent will take a uevent and attempt to match it against the devices classes. If it matches, it will be removed from the device class.
func removeDeviceFromUEvent(udevEvent *udev.UEvent, deviceClasses map[string]*dsv1.DeviceClass) {
	devicePath := prefixSysPath(udevEvent.SysPath)
	for _, devClass := range deviceClasses {
		devClass.RemoveDevice(devicePath)
	}
}