package displayctl

import (
	"context"
	"fmt"

	"edge-infra.dev/pkg/sds/display/displaymanager/manager"
	v2 "edge-infra.dev/pkg/sds/display/k8s/apis/v2"
	"edge-infra.dev/pkg/sds/lib/set"
)

// Upgrades the NodeDisplayConfig, if required.
//
// The NodeDisplayConfig requires an upgrade if either:
//   - It has just been converted by the webhook.
//   - It was previously upgraded, but one or more of the displays
//     were disconnected and couldn't be identified.
func (c *NodeDisplayConfigController) upgradeNodeDisplayConfig(ctx context.Context, nodeDisplayConfig *v2.NodeDisplayConfig) (bool, error) {
	if !requiresUpgrade(nodeDisplayConfig) {
		return false, nil
	}

	// read the current display configuration so we can find which port
	// displays identified by MPID are connected to
	displayConfig, err := readDisplayConfiguration(ctx, c.DisplayManager)
	if err != nil {
		return false, err
	}

	// if the NodeDisplayConfig has V1 spec, it has just been converted
	// by the conversion webhook and needs upgrading
	if nodeDisplayConfig.Spec.HasV1() {
		upgradeV1Spec(nodeDisplayConfig, displayConfig)
		return true, nil
	}

	// if there is no V1 spec, attempt to identify displays which were
	// disconnected during the upgrade
	return identifyDisconnectedDisplays(nodeDisplayConfig, displayConfig), nil
}

func requiresUpgrade(nodeDisplayConfig *v2.NodeDisplayConfig) bool {
	if nodeDisplayConfig.Spec == nil {
		return false
	}

	if nodeDisplayConfig.Spec.HasV1() {
		return true
	}

	return nodeDisplayConfig.Spec.HasDisconnected()
}

func readDisplayConfiguration(ctx context.Context, displayManager manager.DisplayManager) (*v2.DisplayConfig, error) {
	// wait for display manager to become ready before applying configuration
	if err := displayManager.Wait(ctx); err != nil {
		return nil, fmt.Errorf("display manager did not become ready: %w", err)
	}

	// read displays and input devices from node
	displayConfig, _, err := displayManager.Read(ctx)
	if err != nil {
		return nil, fmt.Errorf("unable to read display configuration from node: %w", err)
	}

	return displayConfig, nil
}

func upgradeV1Spec(nodeDisplayConfig *v2.NodeDisplayConfig, displayConfig *v2.DisplayConfig) {
	result, disconnectedIDs := v2.DisplayConfigFromV1(nodeDisplayConfig.Spec.V1, displayConfig)
	nodeDisplayConfig.Spec = result
	if len(disconnectedIDs) > 0 {
		nodeDisplayConfig.Status.Upgrade.DisconnectedDisplays = disconnectedIDs
	}
}

func identifyDisconnectedDisplays(nodeDisplayConfig *v2.NodeDisplayConfig, displayConfig *v2.DisplayConfig) bool {
	identified := set.Set[v2.DisplayPort]{}
	disconnectedIDs := nodeDisplayConfig.DisconnectedDisplayIDs()

	// identify disconnected in displays
	for idx, display := range nodeDisplayConfig.Spec.Displays {
		if mpid, ok := disconnectedIDs[display.DisplayPort]; ok {
			if foundDisplay := displayConfig.Displays.FindByMPID(mpid); foundDisplay != nil {
				identified.Add(display.DisplayPort)
				display.DisplayPort = foundDisplay.DisplayPort
				nodeDisplayConfig.Spec.Displays[idx] = display
			}
		}
	}

	// identify disconnected in layout
	for idx, dp := range nodeDisplayConfig.Spec.Layout {
		if mpid, ok := disconnectedIDs[dp]; ok {
			if foundDisplay := displayConfig.Displays.FindByMPID(mpid); foundDisplay != nil {
				identified.Add(dp)
				nodeDisplayConfig.Spec.Layout[idx] = foundDisplay.DisplayPort
			}
		}
	}

	// updated disconnected displays
	disconnectedIDs = removeIdentifiedFromDisconnectedIDs(disconnectedIDs, identified.ToSlice())
	nodeDisplayConfig.Status.Upgrade.DisconnectedDisplays = disconnectedIDs

	return len(identified) > 0
}

func removeIdentifiedFromDisconnectedIDs(disconnectedIDs map[v2.DisplayPort]v2.MPID, identified []v2.DisplayPort) map[v2.DisplayPort]v2.MPID {
	for _, dp := range identified {
		delete(disconnectedIDs, dp)
	}

	if len(disconnectedIDs) == 0 {
		return nil
	}

	return disconnectedIDs
}