// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

// Package controller implements the node-local control loop inside snapshot-agent.
// It does not own CRDs or replace the operator. Instead it watches pod and work
// order state on the current node and delegates CRIU/CUDA execution to the
// snapshot executor workflows.
package controller

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/go-logr/logr"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	ktypes "k8s.io/apimachinery/pkg/types"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	corev1apply "k8s.io/client-go/applyconfigurations/core/v1"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/dynamic/dynamicinformer"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/workqueue"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/ai-dynamo/snapshot/agent/internal/executor"
	"github.com/ai-dynamo/snapshot/agent/internal/lifecyclehooks"
	"github.com/ai-dynamo/snapshot/agent/internal/nsmount"
	"github.com/ai-dynamo/snapshot/agent/internal/pagebroker"
	snapshotruntime "github.com/ai-dynamo/snapshot/agent/internal/runtime"
	"github.com/ai-dynamo/snapshot/agent/internal/types"
	"github.com/ai-dynamo/snapshot/api/compat"
	"github.com/ai-dynamo/snapshot/api/podcontract"
	snapshotv1alpha1 "github.com/ai-dynamo/snapshot/api/v1alpha1"
)

// NodeController watches local-node pods with checkpoint metadata and reconciles
// snapshot execution for checkpoint and restore requests. Both paths are workqueue
// driven, with typed reads/writes via an uncached controller-runtime client.
type NodeController struct {
	config                    *types.AgentConfig
	clientset                 kubernetes.Interface
	client                    client.Client
	dynClient                 dynamic.Interface
	runtime                   snapshotruntime.Runtime
	injector                  executor.RestoreMounter
	log                       logr.Logger
	checkpointFn              func(ctx context.Context, params CheckpointParams) (*pagebroker.PublishedArtifact, error)
	fetchManifestFn           func(ctx context.Context, artifact *pagebroker.PublishedArtifact) (*types.CheckpointManifest, error)
	restoreFn                 func(context.Context, snapshotruntime.Runtime, logr.Logger, executor.RestoreRequest, executor.RestoreMounter) (executor.RestoreResult, error)
	writeControlSentinelFn    func(int, string, []byte) error
	writePodControlSentinelFn func(string, string, string, []byte) error
	controlSentinelExistsFn   func(int, string) (bool, error)
	sendSignalFn              func(logr.Logger, int, syscall.Signal, string) error
	stagingCapacityFn         func(string) (int64, error)
	minimumImageSizeFn        func(string, []int) (int64, error)
	restoreQueue              workqueue.TypedDelayingInterface[client.ObjectKey]
	restorePodLister          corev1listers.PodLister
	compareFn                 func(compat.Gate, compat.Environment, compat.Environment) []compat.Mismatch
	withoutProxy              bool

	// captureQueue holds PodSnapshotContent names. Its per-key exclusion is the capture path's only
	// lock, which is enough because exactly one agent runs per node.
	captureQueue workqueue.TypedRateLimitingInterface[string]

	// checkpointGates holds a *checkpointGate per work order, while the agent waits for its
	// source container to be ready for the checkpoint.
	checkpointGates sync.Map
	// hookExec runs a lifecycle hook's file action in its container, and
	// hookHTTP sends its HTTP requests. Tests replace them.
	hookExec lifecyclehooks.Execer
	hookHTTP *http.Client
	// nowFn and requeueCaptureFn are time.Now and captureQueue.AddAfter. Tests
	// replace them.
	nowFn            func() time.Time
	requeueCaptureFn func(name string, after time.Duration)

	// postRestoreWaitFn waits between post-restore attempts: waitOrDone.
	// Tests replace it.
	postRestoreWaitFn func(ctx context.Context, d time.Duration) error

	handledRestores sync.Map

	// startedRestoreContainers holds, per restoreContainerKey, the ID of a
	// restore destination container the NRI plugin saw start, so a restore
	// begins before kubelet publishes the container in the Pod status.
	startedRestoreContainers sync.Map

	// contentIndexer is the PodSnapshotContent informer's indexer, indexed by source pod
	// (podRefIndex). The source-pod informer uses it to map a pod event back to its work order.
	contentIndexer cache.Indexer

	stopCh chan struct{}
}

type restoreArtifact struct {
	SnapshotName        string
	ContentUID          string
	SourceContainerName string
	Path                string
	// PublishedArtifact is set only for bound content, read from status.
	PublishedArtifact *pagebroker.PublishedArtifact
}

type restoreTarget struct {
	SnapshotName        string
	ContentUID          string
	SourceContainerName string
}

type restorePlan struct {
	artifact *restoreArtifact
	mappings []podcontract.ContainerMapping
	// Read once in preflight, so both gates reach the same answer.
	skipCompatCheck bool
	// postRestore runs the checkpoint's lifecycle.postRestore action.
	postRestore postRestoreHandler
}

type restoreResultState int

const (
	restoreResultPending restoreResultState = iota
	restoreResultSucceeded
	restoreResultFailed
	restoreResultIncompatible
)

type restoreResult struct {
	destination string
	state       restoreResultState
	// reason explains a failed or incompatible outcome.
	reason string
}

type restorePendingError struct {
	reason  string
	message string
}

func (e *restorePendingError) Error() string {
	return e.message
}

type restoreOperation struct {
	controller      *NodeController
	pod             *corev1.Pod
	artifact        *restoreArtifact
	postRestore     postRestoreHandler
	skipCompatCheck bool
	destination     string
	containerID     string
	startedAt       time.Time
	log             logr.Logger
}

const (
	containerResolveAttemptTimeout     = 1 * time.Second
	restoreContainerResolveInterval    = 50 * time.Millisecond
	restoreContainerResolveTimeout     = 30 * time.Second
	restorePreflightManifestTimeout    = 30 * time.Second
	restoreAlreadyCompletedReason      = "RestoreAlreadyCompleted"
	restoreAlreadyPartialReason        = "RestoreAlreadyPartiallySucceeded"
	restoreAlreadyFailedReason         = "RestoreAlreadyFailed"
	restoreFinalizerUpdateFailedReason = "RestoreFinalizerUpdateFailed"
	restoreStatusUpdateFailedReason    = "RestoreStatusUpdateFailed"
	restoreRequestedReason             = "RestoreRequested"
	restoreCompatUncheckedReason       = "RestoreCompatibilityUnchecked"
	restoreStatusFieldManager          = "snapshot-agent-restore"
	restorePodFinalizer                = "snapshot/restore-protection"
	snapshotEventComponent             = "snapshot"
	restoreSafetyRequeueInterval       = 30 * time.Second

	// snapshotContentResyncInterval re-drives every PodSnapshotContent work order so a
	// not-yet-Ready source pod is re-checked for quiesce without a busy loop.
	snapshotContentResyncInterval = 10 * time.Second

	// waitingForCheckpointReadinessReason is the Event on the source Pod when
	// the agent starts waiting for its checkpoint readiness probe.
	waitingForCheckpointReadinessReason = "WaitingForCheckpointReadiness"
)

// podSnapshotContentGVR is the cluster-scoped resource the capture informer watches.
var podSnapshotContentGVR = snapshotv1alpha1.GroupVersion.WithResource("podsnapshotcontents")

// restoreFailedSentinelContents is informational: readers of restore-failed
// act on its existence only.
var restoreFailedSentinelContents = []byte("failed\n")

// NewNodeController creates the node-local controller that runs inside snapshot-agent.
func NewNodeController(
	cfg *types.AgentConfig,
	rt snapshotruntime.Runtime,
	log logr.Logger,
) (*NodeController, error) {
	restConfig, err := rest.InClusterConfig()
	if err != nil {
		return nil, fmt.Errorf("failed to get in-cluster config: %w", err)
	}

	clientset, err := kubernetes.NewForConfig(restConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create kubernetes client: %w", err)
	}

	scheme := runtime.NewScheme()
	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
	utilruntime.Must(snapshotv1alpha1.AddToScheme(scheme))

	typedClient, err := client.New(restConfig, client.Options{Scheme: scheme})
	if err != nil {
		return nil, fmt.Errorf("failed to create typed client: %w", err)
	}

	dynClient, err := dynamic.NewForConfig(restConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create dynamic client: %w", err)
	}

	nsm := nsmount.New(log)
	return newDefaultController(cfg, clientset, typedClient, dynClient, rt, nsm, log), nil
}

func newDefaultController(
	cfg *types.AgentConfig,
	clientset kubernetes.Interface,
	typedClient client.Client,
	dynClient dynamic.Interface,
	rt snapshotruntime.Runtime,
	injector executor.RestoreMounter,
	log logr.Logger,
) *NodeController {
	w := &NodeController{
		config:    cfg,
		clientset: clientset,
		client:    typedClient,
		dynClient: dynClient,
		runtime:   rt,
		injector:  injector,
		log:       log,
		stopCh:    make(chan struct{}),
		restoreQueue: workqueue.NewTypedDelayingQueueWithConfig(
			workqueue.TypedDelayingQueueConfig[client.ObjectKey]{Name: "restore-pods"},
		),
		captureQueue: workqueue.NewTypedRateLimitingQueueWithConfig(
			workqueue.DefaultTypedControllerRateLimiter[string](),
			workqueue.TypedRateLimitingQueueConfig[string]{Name: "capture-contents"},
		),

		restoreFn:                 executor.Restore,
		writeControlSentinelFn:    snapshotruntime.WriteControlSentinel,
		writePodControlSentinelFn: snapshotruntime.WritePodControlSentinel,
		controlSentinelExistsFn:   snapshotruntime.ControlSentinelExists,
		sendSignalFn:              snapshotruntime.SendSignalToPID,
		stagingCapacityFn:         snapshotruntime.StagingCapacity,
		minimumImageSizeFn:        snapshotruntime.MinimumImageSize,
		compareFn:                 compat.Compare,
		hookExec:                  rt,
		hookHTTP:                  lifecyclehooks.NewHTTPClient(),
		nowFn:                     time.Now,
		postRestoreWaitFn:         waitOrDone,
	}
	w.requeueCaptureFn = w.captureQueue.AddAfter
	w.checkpointFn = w.executorCheckpoint
	w.fetchManifestFn = func(ctx context.Context, artifact *pagebroker.PublishedArtifact) (*types.CheckpointManifest, error) {
		return executor.FetchArtifactManifest(ctx, cfg.PageBroker.ControlSocketPath, artifact)
	}
	return w
}

// RestoreWithoutProxy makes restores run without the restore proxy, on a node
// whose runtime cannot run it. Call it before Run.
func (w *NodeController) RestoreWithoutProxy() {
	w.withoutProxy = true
}

// Run starts the local pod informers and processes checkpoint/restore events.
func (w *NodeController) Run(ctx context.Context) error {
	defer w.restoreQueue.ShutDown()
	defer w.captureQueue.ShutDown()
	// Seed the agent logger onto ctx so the capture path resolves it via log.FromContext.
	ctx = logr.NewContext(ctx, w.log)
	w.log.Info("Starting snapshot node controller",
		"node", w.config.NodeName,
		"restore_from_annotation", podcontract.RestoreFromAnnotation,
	)

	w.log.Info("Watching pods cluster-wide (all namespaces)")

	var syncFuncs []cache.InformerSynced

	restoreFactoryOpts := []informers.SharedInformerOption{
		informers.WithTweakListOptions(tweakNodePodListOptions(w.config.NodeName)),
	}

	restoreFactory := informers.NewSharedInformerFactoryWithOptions(
		w.clientset, 30*time.Second, restoreFactoryOpts...,
	)

	restorePods := restoreFactory.Core().V1().Pods()
	restoreInformer := restorePods.Informer()
	w.restorePodLister = restorePods.Lister()
	if _, err := restoreInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: w.enqueueRestorePod,
		UpdateFunc: func(_, newObj interface{}) {
			w.enqueueRestorePod(newObj)
		},
		DeleteFunc: w.forgetRestorePod,
	}); err != nil {
		return fmt.Errorf("failed to add restore informer handler: %w", err)
	}
	go restoreFactory.Start(w.stopCh)
	syncFuncs = append(syncFuncs, restoreInformer.HasSynced)

	// Capture path: a dynamic informer over PodSnapshotContent work orders, filtered at
	// the list/watch level to this node's mirror label. The node-label filter is the
	// node scoping; reconcileCapture keeps a defensive nodeName check.
	nodeContentSelector := labels.SelectorFromSet(labels.Set{snapshotv1alpha1.SnapshotNodeLabel: w.config.NodeName}).String()
	dynFactory := dynamicinformer.NewFilteredDynamicSharedInformerFactory(
		w.dynClient, snapshotContentResyncInterval, metav1.NamespaceAll,
		func(opts *metav1.ListOptions) {
			opts.LabelSelector = nodeContentSelector
		},
	)
	contentInformer := dynFactory.ForResource(podSnapshotContentGVR).Informer()
	// Index work orders by their source pod so a source-pod event maps back to its
	// PodSnapshotContent in O(1). Must be registered before the informer starts.
	if err := contentInformer.AddIndexers(cache.Indexers{podRefIndex: podRefIndexFunc}); err != nil {
		return fmt.Errorf("failed to add snapshot-content podRef indexer: %w", err)
	}
	w.contentIndexer = contentInformer.GetIndexer()
	// Handlers only enqueue: reconciling inline would block delivery for the length of a dump.
	// The resync re-enqueues everything, which is the backstop for a not-yet-quiesced source.
	if _, err := contentInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: w.enqueueContent,
		UpdateFunc: func(_, newObj interface{}) {
			w.enqueueContent(newObj)
		},
	}); err != nil {
		return fmt.Errorf("failed to add snapshot-content informer handler: %w", err)
	}
	go dynFactory.Start(w.stopCh)
	syncFuncs = append(syncFuncs, contentInformer.HasSynced)

	// Source-pod informer: keyed on CaptureEligibleLabel, the promotion label reconcileCapture adds
	// only after a source pod passes validation, so only validated pods feed the queue.
	// A pod status change (a checkpoint container crashing, or the target becoming ready) does
	// not touch the PodSnapshotContent, so without this trigger it would only be acted on at the
	// content informer's resync. It needs its own factory: its selector is disjoint from the restore
	// informer's.
	sourceSelector := labels.SelectorFromSet(labels.Set{snapshotv1alpha1.CaptureEligibleLabel: "true"}).String()
	sourceFactoryOpts := []informers.SharedInformerOption{
		informers.WithTweakListOptions(tweakLabeledNodePodListOptions(sourceSelector, w.config.NodeName)),
	}
	sourceFactory := informers.NewSharedInformerFactoryWithOptions(
		w.clientset, 30*time.Second, sourceFactoryOpts...,
	)
	sourceInformer := sourceFactory.Core().V1().Pods().Informer()
	if _, err := sourceInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: w.enqueueCaptureForSourcePod,
		UpdateFunc: func(_, newObj interface{}) {
			w.enqueueCaptureForSourcePod(newObj)
		},
	}); err != nil {
		return fmt.Errorf("failed to add source-pod informer handler: %w", err)
	}
	go sourceFactory.Start(w.stopCh)
	syncFuncs = append(syncFuncs, sourceInformer.HasSynced)

	// Close stopCh on cancellation so a stalled cache sync (below) is unblocked by ctx, not only on
	// the happy path.
	var stopOnce sync.Once
	go func() {
		<-ctx.Done()
		w.restoreQueue.ShutDown()
		w.captureQueue.ShutDown()
		stopOnce.Do(func() { close(w.stopCh) })
	}()

	if !cache.WaitForCacheSync(w.stopCh, syncFuncs...) {
		return fmt.Errorf("failed to sync informer caches")
	}

	go w.runRestoreQueue(ctx)
	go w.runCaptureQueue(ctx)
	w.log.Info("PodSnapshot node controller started and caches synced")
	<-ctx.Done()
	stopOnce.Do(func() { close(w.stopCh) })
	return nil
}

func tweakNodePodListOptions(nodeName string) func(*metav1.ListOptions) {
	return func(opts *metav1.ListOptions) {
		opts.FieldSelector = fields.OneTermEqualSelector("spec.nodeName", nodeName).String()
	}
}

func tweakLabeledNodePodListOptions(labelSelector, nodeName string) func(*metav1.ListOptions) {
	return func(opts *metav1.ListOptions) {
		tweakNodePodListOptions(nodeName)(opts)
		opts.LabelSelector = labelSelector
	}
}

func (w *NodeController) enqueueRestorePod(obj interface{}) {
	pod, ok := podFromInformerObj(obj)
	if !ok || !w.restorePodRelevant(pod) {
		return
	}
	w.restoreQueue.Add(client.ObjectKeyFromObject(pod))
}

func (w *NodeController) forgetRestorePod(obj interface{}) {
	pod, ok := podFromInformerObj(obj)
	if !ok {
		return
	}
	w.handledRestores.Delete(string(pod.UID))
	prefix := string(pod.UID) + "/"
	w.startedRestoreContainers.Range(func(key, _ any) bool {
		if strings.HasPrefix(key.(string), prefix) {
			w.startedRestoreContainers.Delete(key)
		}
		return true
	})
}

// RestoreContainerStarted records a restore destination's container that the
// NRI plugin saw start, and requeues its Pod so the restore begins without
// waiting for kubelet to publish the container in the Pod status. In a
// vCluster tenant that status is also the only other way to find the
// container: the runtime labels it with the host Pod's name.
func (w *NodeController) RestoreContainerStarted(pod ktypes.NamespacedName, uid ktypes.UID, container, containerID string) {
	w.startedRestoreContainers.Store(restoreContainerKey(uid, container), containerID)
	w.restoreQueue.Add(pod)
}

// RestoreContainerStopped forgets a reported destination container that has
// stopped, unless a later container already replaced the report. A restore
// into it would fail the Pod's restore instead of waiting for the container
// kubelet starts next.
func (w *NodeController) RestoreContainerStopped(uid ktypes.UID, container, containerID string) {
	w.startedRestoreContainers.CompareAndDelete(restoreContainerKey(uid, container), containerID)
}

func restoreContainerKey(uid ktypes.UID, container string) string {
	return string(uid) + "/" + container
}

func (w *NodeController) restorePodRequested(pod *corev1.Pod) bool {
	if pod.Spec.NodeName != w.config.NodeName {
		return false
	}
	_, requested := pod.Annotations[podcontract.RestoreFromAnnotation]
	return requested
}

func (w *NodeController) restorePodRelevant(pod *corev1.Pod) bool {
	return pod.Spec.NodeName == w.config.NodeName &&
		(hasFinalizer(pod, restorePodFinalizer) || w.restorePodRequested(pod))
}

// runRestoreQueue gives each pod its own goroutine. A CRIU run holds it for minutes and must not
// head-of-line block unrelated pods, and there is no fixed worker count to bound it: the node's
// GPUs and pod capacity already cap how many restores can exist, and a count low enough to matter
// would instead leave deleted pods Terminating behind the queue, since the restore finalizer comes
// off in processRestoreQueueItem.
func (w *NodeController) runRestoreQueue(ctx context.Context) {
	for {
		key, shutdown := w.restoreQueue.Get()
		if shutdown {
			return
		}
		go w.processRestoreQueueItem(ctx, key)
	}
}

func (w *NodeController) processRestoreQueueItem(ctx context.Context, key client.ObjectKey) {
	requeue := false
	defer func() {
		w.restoreQueue.Done(key)
		if requeue && ctx.Err() == nil {
			w.restoreQueue.AddAfter(key, restoreSafetyRequeueInterval)
		}
	}()
	// Registered after the Done defer so it runs first, leaving requeue false.
	var panicked *corev1.Pod
	defer recoverPanic(w.log, "restore work order", func(err error) {
		w.failRestoreAfterPanic(ctx, key, panicked, err)
	}, "pod", key.String())

	pod, err := w.restorePodLister.Pods(key.Namespace).Get(key.Name)
	if apierrors.IsNotFound(err) {
		return
	}
	if err != nil {
		w.log.Error(err, "Failed to read restore pod from informer cache", "pod", key.String())
		requeue = true
		return
	}
	pod = pod.DeepCopy()
	panicked = pod
	if w.skipRequestedAfterRefusal(pod) {
		// The skip request is the way back for a pod the gates turned down, so
		// it has to clear the in-process marker as well as the condition below.
		w.handledRestores.Delete(string(pod.UID))
	} else if w.restoreHandled(pod) {
		requeue = w.removeRestoreFinalizerWithEvent(ctx, pod)
		return
	} else if isRestoreTerminal(pod) {
		requeue = w.handleTerminalRestorePod(ctx, pod)
		return
	}
	if pod.DeletionTimestamp != nil {
		requeue = w.removeRestoreFinalizerWithEvent(ctx, pod)
		return
	}
	if !isRestorePodActive(pod) {
		requeue = w.failRestorePod(ctx, pod, fmt.Errorf("Pod %s/%s is not eligible for restore in phase %s; expected Pending or Running", pod.Namespace, pod.Name, pod.Status.Phase))
		return
	}
	requeue = w.reconcileRestorePod(ctx, pod)
}

// handleTerminalRestorePod reports the previously recorded terminal outcome
// and removes restore protection without replaying any worker.
func (w *NodeController) handleTerminalRestorePod(ctx context.Context, pod *corev1.Pod) bool {
	condition := findRestoredCondition(pod)
	eventType := corev1.EventTypeWarning
	reason := restoreAlreadyFailedReason
	message := condition.Message

	switch {
	case isRestoreSucceeded(pod):
		eventType = corev1.EventTypeNormal
		reason = restoreAlreadyCompletedReason
		message = "Pod restore already completed; no further action is required"
	case isRestorePartiallySucceeded(pod):
		reason = restoreAlreadyPartialReason
		if message == "" {
			message = "Pod restore partially succeeded; create a new restore Pod to retry failed destinations"
		}
	case message == "":
		message = "Pod restore previously failed; create a new restore Pod to retry"
	}

	// finishRestore's marker is in-memory, so an agent restart sends every
	// terminal restore pod here; without marking again the resync reports it
	// once per interval for as long as the pod exists.
	w.markRestoreHandled(pod)
	emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, eventType, reason, message)
	return w.removeRestoreFinalizerWithEvent(ctx, pod)
}

func (w *NodeController) reconcileRestorePod(ctx context.Context, pod *corev1.Pod) bool {
	podKey := fmt.Sprintf("%s/%s", pod.Namespace, pod.Name)
	plan, err := w.preflightRestore(ctx, pod)
	if err != nil {
		return w.handleRestorePreflightError(ctx, pod, err)
	}
	if err := w.addRestoreFinalizer(ctx, pod); err != nil {
		w.log.Error(err, "Failed to protect restore Pod", "pod", podKey)
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreFinalizerUpdateFailedReason, err.Error())
		return true
	}
	return w.restorePodContainers(ctx, pod, plan, podKey)
}

// preflightRestore validates the restore Pod, its referenced PodSnapshot and
// bound PodSnapshotContent, and the node-local artifact. Restore execution is
// entered only when this method returns a complete artifact.
func (w *NodeController) preflightRestore(ctx context.Context, pod *corev1.Pod) (*restorePlan, error) {
	snapshot, err := w.getPodSnapshotFromPod(ctx, pod)
	if err != nil {
		return nil, err
	}
	if err := validatePodSnapshotForRestore(snapshot); err != nil {
		return nil, err
	}
	content, err := w.getPodSnapshotContentFromSnapshot(ctx, pod, snapshot)
	if err != nil {
		return nil, err
	}
	if err := validatePodSnapshotContentForRestore(content); err != nil {
		return nil, err
	}
	target, mappings, err := validateRestoreTarget(pod, snapshot, content)
	if err != nil {
		return nil, err
	}
	artifact, err := w.resolveRestoreArtifact(fmt.Sprintf("%s/%s", pod.Namespace, pod.Name), target, content)
	if err != nil {
		return nil, err
	}
	// Gate A: the earliest point the checkpoint's own record of what it was
	// captured on is readable, and still before any of the restore is attempted.
	skipCompatCheck := w.skipCompatCheckRequested(pod)
	if err := w.preflightCompatibility(ctx, pod, artifact, mappings, skipCompatCheck); err != nil {
		return nil, err
	}
	if w.config.CRIU.TcpEstablished && pod.Status.PodIP == "" {
		return nil, newRestorePendingError("PodIPPending", fmt.Sprintf("Waiting for restore Pod %s/%s to receive an IP address", pod.Namespace, pod.Name))
	}
	return &restorePlan{
		artifact:        artifact,
		mappings:        mappings,
		skipCompatCheck: skipCompatCheck,
		postRestore:     w.newPostRestoreHandler(postRestoreAction(content, target.SourceContainerName)),
	}, nil
}

func (w *NodeController) getPodSnapshotFromPod(ctx context.Context, pod *corev1.Pod) (*snapshotv1alpha1.PodSnapshot, error) {
	snapshotName, err := podcontract.GetRestoreFromSnapshotName(pod.Annotations)
	if err != nil {
		return nil, err
	}
	snapshot := &snapshotv1alpha1.PodSnapshot{}
	key := client.ObjectKey{Namespace: pod.Namespace, Name: snapshotName}
	if err := w.client.Get(ctx, key, snapshot); err != nil {
		if apierrors.IsNotFound(err) {
			if restoreInProgress(pod) {
				return nil, fmt.Errorf("PodSnapshot %s disappeared while restore was in progress", key.String())
			}
			return nil, newRestorePendingError("SnapshotPending", fmt.Sprintf("Waiting for PodSnapshot %s", key.String()))
		}
		w.log.Error(err, "Failed to read PodSnapshot during restore preflight", "pod", client.ObjectKeyFromObject(pod).String(), "snapshot", key.String())
		return nil, newRestorePendingError("SnapshotPending", fmt.Sprintf("Unable to read PodSnapshot %s; retrying", key.String()))
	}
	return snapshot, nil
}

func validatePodSnapshotForRestore(snapshot *snapshotv1alpha1.PodSnapshot) error {
	key := client.ObjectKeyFromObject(snapshot)
	if snapshot.DeletionTimestamp != nil {
		return fmt.Errorf("PodSnapshot %s is deleting", key.String())
	}
	if snapshotv1alpha1.IsPodSnapshotFailed(snapshot) {
		cond := meta.FindStatusCondition(snapshot.Status.Conditions, snapshotv1alpha1.PodSnapshotConditionFailed)
		message := fmt.Sprintf("PodSnapshot %s failed", key.String())
		if cond != nil && cond.Message != "" {
			message += ": " + cond.Message
		}
		return errors.New(message)
	}
	if !snapshotv1alpha1.IsPodSnapshotSucceeded(snapshot) || snapshot.Status.BoundPodSnapshotContentName == nil ||
		strings.TrimSpace(*snapshot.Status.BoundPodSnapshotContentName) == "" {
		return newRestorePendingError("SnapshotPending", fmt.Sprintf("Waiting for PodSnapshot %s to become Ready", key.String()))
	}
	return nil
}

func (w *NodeController) getPodSnapshotContentFromSnapshot(ctx context.Context, pod *corev1.Pod, snapshot *snapshotv1alpha1.PodSnapshot) (*snapshotv1alpha1.PodSnapshotContent, error) {
	contentName := strings.TrimSpace(ptr.Deref(snapshot.Status.BoundPodSnapshotContentName, ""))
	content := &snapshotv1alpha1.PodSnapshotContent{}
	if err := w.client.Get(ctx, client.ObjectKey{Name: contentName}, content); err != nil {
		if apierrors.IsNotFound(err) {
			if restoreInProgress(pod) {
				return nil, fmt.Errorf("bound PodSnapshotContent %s disappeared while restore was in progress", contentName)
			}
			return nil, newRestorePendingError("ContentPending", fmt.Sprintf("Waiting for bound PodSnapshotContent %s", contentName))
		}
		w.log.Error(err, "Failed to read PodSnapshotContent during restore preflight", "snapshot", client.ObjectKeyFromObject(snapshot).String(), "content", contentName)
		return nil, newRestorePendingError("ContentPending", fmt.Sprintf("Unable to read bound PodSnapshotContent %s; retrying", contentName))
	}
	ref := content.Spec.PodSnapshotRef
	if ref.Namespace != snapshot.Namespace || ref.Name != snapshot.Name || ref.UID == "" || ref.UID != snapshot.UID {
		return nil, fmt.Errorf("PodSnapshotContent %s has a stale backlink to %s/%s uid %q", content.Name, ref.Namespace, ref.Name, ref.UID)
	}
	return content, nil
}

func validatePodSnapshotContentForRestore(content *snapshotv1alpha1.PodSnapshotContent) error {
	if content.DeletionTimestamp != nil {
		return fmt.Errorf("PodSnapshotContent %s is deleting", content.Name)
	}
	if snapshotv1alpha1.IsPodSnapshotContentFailed(content) {
		cond := meta.FindStatusCondition(content.Status.Conditions, snapshotv1alpha1.PodSnapshotConditionFailed)
		message := fmt.Sprintf("PodSnapshotContent %s failed", content.Name)
		if cond != nil && cond.Message != "" {
			message += ": " + cond.Message
		}
		return errors.New(message)
	}
	if !snapshotv1alpha1.IsPodSnapshotContentSucceeded(content) {
		return newRestorePendingError("ContentPending", fmt.Sprintf("Waiting for PodSnapshotContent %s to become Ready", content.Name))
	}
	return nil
}

// RestoreDestinations returns the names of the containers a restore Pod
// restores into. When the Pod carries no container map, the destination is the
// captured container, which the PodSnapshot's spec names. That spec is set at
// creation and immutable, so this works while the snapshot is still being
// taken. It checks nothing else: the restore itself still runs the full
// preflight, including the bound PodSnapshotContent.
func (w *NodeController) RestoreDestinations(ctx context.Context, pod *corev1.Pod) ([]string, error) {
	source := ""
	if _, mapped := pod.Annotations[podcontract.RestoreContainerMapAnnotation]; !mapped {
		snapshot, err := w.getPodSnapshotFromPod(ctx, pod)
		if err != nil {
			return nil, err
		}
		if source, err = singleTargetContainer(snapshot.Spec.Source.PodRef.Containers); err != nil {
			return nil, err
		}
	}
	mappings, err := podcontract.ContainerMappingsFromAnnotations(pod.Annotations, source)
	if err != nil {
		return nil, err
	}
	destinations := make([]string, 0, len(mappings))
	for _, mapping := range mappings {
		destinations = append(destinations, mapping.Destination)
	}
	return destinations, nil
}

// validateRestoreTarget resolves the captured source and validates every
// requested destination against the restore Pod.
func validateRestoreTarget(pod *corev1.Pod, snapshot *snapshotv1alpha1.PodSnapshot, content *snapshotv1alpha1.PodSnapshotContent) (*restoreTarget, []podcontract.ContainerMapping, error) {
	containerName, err := singleTargetContainer(content.Spec.Source.PodRef.Containers)
	if err != nil {
		return nil, nil, err
	}
	mappings, err := podcontract.ContainerMappingsFromAnnotations(pod.Annotations, containerName)
	if err != nil {
		return nil, nil, err
	}
	if err := podcontract.Validate(pod, podcontract.Request{
		SnapshotName:    snapshot.Name,
		SourceContainer: containerName,
		Mappings:        mappings,
	}); err != nil {
		return nil, nil, err
	}
	return &restoreTarget{SnapshotName: snapshot.Name, ContentUID: string(content.UID), SourceContainerName: containerName}, mappings, nil
}

// resolveRestoreArtifact resolves the validated restore target to its physical
// checkpoint directory. A nil error always returns a complete artifact.
// resolveRestoreArtifact: legacy content is ready when the filesystem artifact
// exists; bound content is ready once status records a published container.
func (w *NodeController) resolveRestoreArtifact(
	podKey string,
	target *restoreTarget,
	content *snapshotv1alpha1.PodSnapshotContent,
) (*restoreArtifact, error) {
	if content.Spec.Storage != nil {
		return w.resolveBoundRestoreArtifact(target, content)
	}
	path, err := nsmount.ResolveArtifactPath(w.config.Storage.BasePath, target.ContentUID, target.SourceContainerName)
	if err != nil {
		return nil, err
	}
	ready, err := w.restoreArtifactReady(w.log, podKey, path)
	if err != nil {
		var pathErr *os.PathError
		if errors.As(err, &pathErr) {
			w.log.Error(err, "Unable to inspect restore artifact", "pod", podKey, "artifact_path", path)
			return nil, newRestorePendingError("ArtifactPending", fmt.Sprintf("Unable to inspect artifact %s; retrying", path))
		}
		return nil, err
	}
	if !ready {
		// Ready guarantees a usable artifact; absence is stale state, not a retryable dependency.
		return nil, fmt.Errorf("artifact %s for Ready PodSnapshotContent %s is missing", path, content.Name)
	}
	return &restoreArtifact{
		SnapshotName:        target.SnapshotName,
		ContentUID:          target.ContentUID,
		SourceContainerName: target.SourceContainerName,
		Path:                path,
	}, nil
}

func (w *NodeController) resolveBoundRestoreArtifact(
	target *restoreTarget,
	content *snapshotv1alpha1.PodSnapshotContent,
) (*restoreArtifact, error) {
	if content.Status.Storage != nil {
		for _, published := range content.Status.Storage.Artifacts {
			if published.ContainerName != target.SourceContainerName {
				continue
			}
			return &restoreArtifact{
				SnapshotName:        target.SnapshotName,
				ContentUID:          target.ContentUID,
				SourceContainerName: target.SourceContainerName,
				PublishedArtifact: &pagebroker.PublishedArtifact{
					StoreId:        content.Spec.Storage.StoreID,
					ArtifactHandle: published.ArtifactHandle,
				},
			}, nil
		}
	}
	return nil, newRestorePendingError("ArtifactPending",
		fmt.Sprintf("Waiting for a published artifact for container %s", target.SourceContainerName))
}

// restorePodContainers owns one restore pass for the Pod. Workers never write
// Pod status; they report one result each and this coordinator publishes the
// aggregate outcome.
func (w *NodeController) restorePodContainers(ctx context.Context, pod *corev1.Pod, plan *restorePlan, podKey string) bool {
	// Read the incoming state before persisting this pass. On agent restart,
	// RestoreInProgress makes each worker check its completion sentinel before
	// considering a CRIU replay.
	recovering := restoreInProgress(pod)
	message := fmt.Sprintf("Restoring %d destination container(s) from PodSnapshot %s", len(plan.mappings), plan.artifact.SnapshotName)
	if err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, podcontract.RestoreReasonInProgress, message); err != nil {
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, err.Error())
		return true
	}

	results := make([]restoreResult, len(plan.mappings))
	var workers sync.WaitGroup
	for i, mapping := range plan.mappings {
		i, destination := i, mapping.Destination
		workers.Go(func() {
			// sync.WaitGroup.Go recovers only to re-panic, so a panic here
			// would kill the agent. Report it as a failed destination instead.
			results[i] = restoreResult{destination: destination, state: restoreResultFailed,
				reason: "agent panicked restoring this container"}
			defer recoverPanic(w.log, "restore destination", nil, "pod", podKey, "destination", destination)
			results[i] = w.restoreDestination(ctx, pod, plan, destination, podKey, recovering)
		})
	}
	workers.Wait()

	return w.recordRestoreResults(ctx, pod, plan.artifact, results)
}

type restoreTally struct {
	total                  int
	succeeded              []string
	failed                 []string
	incompatible           []string
	pending                []string
	incompatibilityReasons []string
	// failureReasons holds "<destination>: <cause>" for each failed destination.
	failureReasons []string
}

type restoreVerdict struct {
	status  corev1.ConditionStatus
	reason  string
	message string
}

// recordRestoreResults publishes the aggregate Pod outcome after every worker
// in the current pass has returned.
func (w *NodeController) recordRestoreResults(ctx context.Context, pod *corev1.Pod, artifact *restoreArtifact, results []restoreResult) bool {
	tally := tallyRestoreResults(results)
	verdict := tally.verdict(artifact.SnapshotName)
	if len(tally.pending) != 0 {
		// The pass is not over, so a write that fails is reported and dropped
		// rather than retried: the next pass publishes again.
		if err := w.applyRestoredCondition(ctx, pod, verdict.status, verdict.reason, verdict.message); err != nil {
			emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, err.Error())
		}
		return true
	}
	return w.finishRestore(ctx, pod, verdict.status, verdict.reason, verdict.message) != nil
}

func tallyRestoreResults(results []restoreResult) restoreTally {
	tally := restoreTally{total: len(results)}
	for _, result := range results {
		switch result.state {
		case restoreResultSucceeded:
			tally.succeeded = append(tally.succeeded, result.destination)
		case restoreResultFailed:
			tally.failed = append(tally.failed, result.destination)
			reason := result.destination
			if result.reason != "" {
				reason = fmt.Sprintf("%s: %s", result.destination, result.reason)
			}
			tally.failureReasons = append(tally.failureReasons, reason)
		case restoreResultPending:
			tally.pending = append(tally.pending, result.destination)
		case restoreResultIncompatible:
			tally.incompatible = append(tally.incompatible, result.destination)
			reason := result.reason
			if len(results) > 1 {
				reason = fmt.Sprintf("%s: %s", result.destination, reason)
			}
			tally.incompatibilityReasons = append(tally.incompatibilityReasons, reason)
		}
	}
	return tally
}

// logWithStep adds the failing step as a structured key when the error names one.
func logWithStep(log logr.Logger, err error) logr.Logger {
	if step := executor.StepOf(err); step != "" {
		return log.WithValues("step", string(step))
	}
	return log
}

// verdict concludes one pass, still in progress or terminal, from the counts
// alone.
func (t restoreTally) verdict(snapshotName string) restoreVerdict {
	switch {
	case len(t.pending) != 0:
		return restoreVerdict{
			status:  corev1.ConditionFalse,
			reason:  podcontract.RestoreReasonInProgress,
			message: t.progressMessage(snapshotName),
		}
	case len(t.failed) == 0 && len(t.incompatible) == 0:
		return restoreVerdict{
			status: corev1.ConditionTrue,
			reason: podcontract.RestoreReasonSucceeded,
			message: fmt.Sprintf("Restored %d destination container(s) from PodSnapshot %s: %s",
				len(t.succeeded), snapshotName, strings.Join(t.succeeded, ", ")),
		}
	case len(t.succeeded) != 0:
		// Causes can contain commas, so entries are separated by semicolons.
		notRestored := append(append([]string{}, t.failureReasons...), t.incompatible...)
		return restoreVerdict{
			status: corev1.ConditionFalse,
			reason: podcontract.RestoreReasonPartiallySucceeded,
			message: fmt.Sprintf("Restored %d of %d destination containers from PodSnapshot %s; not restored: %s",
				len(t.succeeded), t.total, snapshotName, strings.Join(notRestored, "; ")),
		}
	case len(t.failed) == 0:
		return restoreVerdict{
			status:  corev1.ConditionFalse,
			reason:  podcontract.RestoreReasonIncompatible,
			message: refusalMessage(strings.Join(t.incompatibilityReasons, "; ")),
		}
	case len(t.incompatible) != 0:
		return restoreVerdict{
			status: corev1.ConditionFalse,
			reason: podcontract.RestoreReasonFailed,
			message: fmt.Sprintf("Restore failed for %d destination container(s) and refused %d incompatible destination(s) from PodSnapshot %s: %s",
				len(t.failed), len(t.incompatible), snapshotName, strings.Join(t.failureReasons, "; ")),
		}
	default:
		return restoreVerdict{
			status: corev1.ConditionFalse,
			reason: podcontract.RestoreReasonFailed,
			message: fmt.Sprintf("Restore failed for all %d destination container(s) from PodSnapshot %s: %s",
				len(t.failed), snapshotName, strings.Join(t.failureReasons, "; ")),
		}
	}
}

func (t restoreTally) progressMessage(snapshotName string) string {
	counts := []string{
		fmt.Sprintf("%d succeeded", len(t.succeeded)),
		fmt.Sprintf("%d failed", len(t.failed)),
	}
	if len(t.incompatible) != 0 {
		counts = append(counts, fmt.Sprintf("%d incompatible", len(t.incompatible)))
	}
	counts = append(counts, fmt.Sprintf("%d pending", len(t.pending)))

	return fmt.Sprintf("Restore from PodSnapshot %s remains in progress: %s (%s)",
		snapshotName,
		strings.Join(counts, ", "),
		strings.Join(t.pending, ", "),
	)
}

// restoreDestination resolves and restores one destination independently of
// its siblings, returning only its worker-local outcome to the coordinator.
func (w *NodeController) restoreDestination(
	ctx context.Context,
	pod *corev1.Pod,
	plan *restorePlan,
	destination, podKey string,
	recovering bool,
) restoreResult {
	artifact := plan.artifact
	result := restoreResult{destination: destination, state: restoreResultPending}
	containerID, _ := w.resolveRestoreContainerID(ctx, pod, destination, podKey)
	if containerID == "" {
		return result
	}

	startedAt := time.Now()
	log := w.log.WithValues("pod", podKey, "snapshot", artifact.SnapshotName, "destination", destination)
	log.Info("Restore target detected, triggering external restore",
		"content_uid", artifact.ContentUID,
		"source_container", artifact.SourceContainerName,
		"container_id", containerID,
	)
	emitPodEvent(ctx, w.clientset, log, pod, snapshotEventComponent, corev1.EventTypeNormal, restoreRequestedReason, fmt.Sprintf("Restore requested from PodSnapshot %s for destination %s", artifact.SnapshotName, destination))

	if err := w.runRestore(ctx, pod, plan, destination, containerID, startedAt, recovering); err != nil {
		if errors.Is(err, errRestoreInterrupted) {
			// Not an outcome: the next agent resumes the post-restore action.
			return result
		}
		var incompatible *compat.IncompatibleError
		if errors.As(err, &incompatible) {
			result.state = restoreResultIncompatible
			result.reason = compat.Reasons(incompatible.Mismatches)
			w.logRestoreRefusal(pod, incompatible, result.reason)
			return result
		}
		result.state = restoreResultFailed
		result.reason = err.Error()
		logWithStep(log, err).Error(err, "Restore controller worker failed")
		emitPodEvent(ctx, w.clientset, log, pod, snapshotEventComponent, corev1.EventTypeWarning, "RestoreWorkerFailed", err.Error())
		return result
	}
	result.state = restoreResultSucceeded
	return result
}

// resolveRestoreContainerID uses the kubelet-published status fast path, then
// polls inline while this Pod's queue key remains in progress: first for the
// container the NRI plugin saw start, then the node runtime.
func (w *NodeController) resolveRestoreContainerID(ctx context.Context, pod *corev1.Pod, destination, podKey string) (string, bool) {
	if containerID := restoreContainerIDFromStatus(pod, destination); containerID != "" {
		return containerID, false
	}

	w.log.V(1).Info("Restore pod has no running container in Kubernetes status yet; polling node runtime",
		"pod", podKey,
		"container", destination,
	)
	deadlineAt := time.Now().Add(restoreContainerResolveTimeout)
	deadline := time.NewTimer(time.Until(deadlineAt))
	defer deadline.Stop()
	tick := time.NewTicker(restoreContainerResolveInterval)
	defer tick.Stop()
	for {
		// Taken, not read: a restore consumes it, and a container started
		// later is reported again.
		if started, ok := w.startedRestoreContainers.LoadAndDelete(restoreContainerKey(pod.UID, destination)); ok {
			w.log.V(1).Info("Resolved restore container started under NRI",
				"pod", podKey,
				"container", destination,
				"container_id", started,
			)
			return started.(string), false
		}

		resolveCtx, cancel := restoreContainerResolveAttemptContext(ctx, deadlineAt)
		containerID, err := w.runtime.ResolveContainerIDByPod(resolveCtx, pod.Name, pod.Namespace, destination)
		cancel()
		if err == nil && containerID != "" {
			w.log.V(1).Info("Resolved restore container via node runtime",
				"pod", podKey,
				"container", destination,
				"container_id", containerID,
			)
			return containerID, false
		}

		select {
		case <-deadline.C:
			w.log.V(1).Info("Timed out polling node runtime for restore container",
				"pod", podKey,
				"container", destination,
			)
			return "", true
		case <-ctx.Done():
			return "", false
		case <-tick.C:
		}
	}
}

// runRestore runs the full restore workflow for one destination container:
//  1. Call executor.Restore (inspect placeholder → nsrestore inside namespace).
//     nsrestore clears any stale restore-complete and restore-failed sentinels
//     on the pod control volume before CRIU, so a prior incarnation cannot
//     release the restored process early.
//  2. Write a restore-complete sentinel holding the restored PID: the
//     CRIU-restored process resumes inside the polling loop that waits on
//     this file, exits quiescence, and resumes the engine. On failure, write
//     restore-failed and kill the placeholder instead.
//  3. With a lifecycle.postRestore action, notify the application within the
//     same deadline. Its file is checked before step 2, once the restored
//     files are in place; a failure here also fails the restore and kills the
//     container.
func (w *NodeController) runRestore(ctx context.Context, pod *corev1.Pod, plan *restorePlan, destination, containerID string, startedAt time.Time, recovering bool) (retErr error) {
	op := w.newRestoreOperation(pod, plan, destination, containerID, startedAt)
	// Without this the placeholder keeps waiting for a restore that panicked.
	defer recoverPanic(op.log, "restore", func(err error) {
		retErr = op.failRestore(ctx, err)
	}, "destination", destination)
	if recovering {
		completed, err := op.recoverCompletedRestore(ctx)
		if err != nil {
			return err
		}
		if completed {
			return op.recoverPostRestore(ctx)
		}
	}

	// The Pod spec is checked before any CRIU work, so an action that can
	// never run restores nothing.
	if err := op.postRestore.check(ctx, op); err != nil {
		return op.failRestore(ctx, err)
	}

	// One deadline covers the restore and its post-restore action.
	restoreCtx := ctx
	if timeout := w.config.Restore.RestoreTimeout(); timeout > 0 {
		var cancel context.CancelFunc
		restoreCtx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}

	restored, err := op.executeRestore(restoreCtx)
	if err != nil {
		var incompatible *compat.IncompatibleError
		if errors.As(err, &incompatible) {
			// The placeholder is left running: killing it restarts the container
			// straight back into the same refusal.
			return incompatible
		}

		var cleanupErr *executor.RestoreCleanupError
		if !errors.As(err, &cleanupErr) {
			return op.failRestore(ctx, err)
		}
		logWithStep(op.log, err).Error(cleanupErr, "Restore completed with cleanup errors")
		emitPodEvent(ctx, w.clientset, op.log, pod, snapshotEventComponent, corev1.EventTypeWarning, "RestoreCleanupFailed", cleanupErr.Error())
	}
	if err := op.postRestore.recheck(ctx, op); err != nil {
		return op.failRestore(ctx, err)
	}
	if err := op.completeRestore(ctx, restored); err != nil {
		return err
	}
	if err := op.postRestore.notify(restoreCtx, op); err != nil {
		return op.failPostRestore(ctx, err)
	}
	return nil
}

// recoverPostRestore finishes a restore that an agent restart interrupted
// after restore-complete: only the post-restore action can be left, and CRIU
// must never run again. Its checks already passed in the interrupted run.
//
// The interrupted run's deadline is gone with it, so the action gets a full
// config.restore.timeoutSeconds again: a restart can extend it by one timeout.
func (op *restoreOperation) recoverPostRestore(ctx context.Context) error {
	postCtx := ctx
	if timeout := op.controller.config.Restore.RestoreTimeout(); timeout > 0 {
		var cancel context.CancelFunc
		postCtx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}
	if err := op.postRestore.notify(postCtx, op); err != nil {
		return op.failPostRestore(ctx, err)
	}
	return nil
}

// recoverCompletedRestore avoids replaying CRIU when the destination-scoped
// completion sentinel proves the operation already finished.
func (op *restoreOperation) recoverCompletedRestore(ctx context.Context) (bool, error) {
	condition := findRestoredCondition(op.pod)
	if condition == nil || condition.Status != corev1.ConditionFalse || condition.Reason != podcontract.RestoreReasonInProgress {
		return false, nil
	}

	hostPID, _, err := op.controller.runtime.ResolveContainer(ctx, op.containerID)
	if err != nil {
		return false, fmt.Errorf("resolve restore container before checking completion sentinel: %w", err)
	}
	exists, err := op.controller.controlSentinelExistsFn(hostPID, podcontract.RestoreCompleteFile)
	if err != nil {
		return false, fmt.Errorf("check restore completion sentinel: %w", err)
	}
	if !exists {
		return false, nil
	}
	return true, nil
}

func (w *NodeController) newRestoreOperation(
	pod *corev1.Pod,
	plan *restorePlan,
	destination string,
	containerID string,
	startedAt time.Time,
) *restoreOperation {
	podKey := fmt.Sprintf("%s/%s", pod.Namespace, pod.Name)
	artifact := plan.artifact
	postRestore := plan.postRestore
	if postRestore == nil {
		postRestore = noPostRestore{}
	}
	return &restoreOperation{
		controller:      w,
		pod:             pod,
		artifact:        artifact,
		postRestore:     postRestore,
		skipCompatCheck: plan.skipCompatCheck,
		destination:     destination,
		containerID:     containerID,
		startedAt:       startedAt,
		log:             w.log.WithValues("pod", podKey, "snapshot", artifact.SnapshotName, "content_uid", artifact.ContentUID, "source_container", artifact.SourceContainerName, "destination_container", destination, "container_id", containerID),
	}
}

func (op *restoreOperation) executeRestore(ctx context.Context) (executor.RestoreResult, error) {
	w := op.controller
	req := executor.RestoreRequest{
		ContentUID:                  op.artifact.ContentUID,
		BasePath:                    w.config.Storage.BasePath,
		ContainerID:                 op.containerID,
		StartedAt:                   op.startedAt,
		PodName:                     op.pod.Name,
		PodNamespace:                op.pod.Namespace,
		TargetPodIP:                 op.pod.Status.PodIP,
		ArtifactContainerName:       op.artifact.SourceContainerName,
		DestinationContainerName:    op.destination,
		SkipCompatCheck:             op.skipCompatCheck,
		WithoutProxy:                w.withoutProxy,
		Clientset:                   w.clientset,
		CustomStorageAvailable:      w.config.CustomStorageAvailable,
		PageBrokerControlSocketPath: w.config.PageBroker.ControlSocketPath,
		PageBrokerRestoreMode:       w.config.PageBroker.RestoreMode,
		PublishedArtifact:           op.artifact.PublishedArtifact,
	}
	return w.restoreFn(ctx, w.runtime, op.log, req, w.injector)
}

func (op *restoreOperation) failRestore(ctx context.Context, restoreErr error) error {
	w := op.controller
	logWithStep(op.log, restoreErr).Error(restoreErr, "External restore failed")
	// Re-resolve because restore may fail before discovering the placeholder PID.
	placeholderHostPID, _, err := w.runtime.ResolveContainer(ctx, op.containerID)
	if err != nil {
		// Nothing can be killed, so the marker is the only way the container's
		// proxy learns that no restore is coming. With no process in the
		// container to write through, it goes through the node's path of the
		// control volume. kubelet names that path after the Pod's UID on the
		// node, which in a vCluster differs from this Pod's; the runtime knows
		// it while the container exists.
		podUID := string(op.pod.UID)
		if nodeUID, uidErr := w.runtime.ResolveContainerPodUID(ctx, op.containerID); uidErr == nil {
			podUID = nodeUID
		} else {
			op.log.V(1).Info("Cannot read the Pod's UID on the node; using its own", "error", uidErr.Error())
		}
		if markErr := w.writePodControlSentinelFn(podUID, op.destination, podcontract.RestoreFailedFile, restoreFailedSentinelContents); markErr != nil {
			op.log.Error(markErr, "Failed to write restore-failed sentinel through the host path")
		}
		return errors.Join(restoreErr, fmt.Errorf("placeholder PID could not be resolved after restore failure: %w", err))
	}
	op.markRestoreFailed(placeholderHostPID)
	if err := op.killPlaceholder(placeholderHostPID, "restore failed"); err != nil {
		return errors.Join(restoreErr, fmt.Errorf("placeholder could not be killed after restore failure: %w", err))
	}
	return restoreErr
}

func (op *restoreOperation) completeRestore(ctx context.Context, restored executor.RestoreResult) error {
	w := op.controller
	placeholderHostPID := restored.PlaceholderHostPID
	// Any PID inside the container mount namespace reaches the control
	// volume through /host/proc/<pid>/root.
	contents := podcontract.FormatRestoreComplete(podcontract.RestoreComplete{PID: restored.RestoredPID})
	if err := w.writeControlSentinelFn(placeholderHostPID, podcontract.RestoreCompleteFile, contents); err != nil {
		op.log.Error(err, "Failed to write restore-complete sentinel")
		op.markRestoreFailed(placeholderHostPID)
		if killErr := op.killPlaceholder(placeholderHostPID, "restore sentinel failed"); killErr != nil {
			return errors.Join(fmt.Errorf("failed to write restore-complete sentinel: %w", err), fmt.Errorf("placeholder could not be killed: %w", killErr))
		}
		return fmt.Errorf("failed to write restore-complete sentinel: %w", err)
	}
	return nil
}

// markRestoreFailed writes the restore-failed sentinel before the caller kills
// the placeholder, so a restarted container sees that this Pod's one restore
// was already used instead of waiting for a restore that will not come. It is
// best effort: the kill must happen whether or not the marker was written.
func (op *restoreOperation) markRestoreFailed(placeholderHostPID int) {
	if err := op.controller.writeControlSentinelFn(placeholderHostPID, podcontract.RestoreFailedFile, restoreFailedSentinelContents); err != nil {
		op.log.Error(err, "Failed to write restore-failed sentinel")
	}
}

// killPlaceholder kills the destination container's init process after
// markRestoreFailed. The restore proxy exits by itself once it sees
// restore-failed, so it may already be gone: that is the goal of the kill,
// not an error.
func (op *restoreOperation) killPlaceholder(placeholderHostPID int, reason string) error {
	err := op.controller.sendSignalFn(op.log, placeholderHostPID, syscall.SIGKILL, reason)
	if errors.Is(err, syscall.ESRCH) {
		op.log.V(1).Info("Placeholder already exited", "pid", placeholderHostPID)
		return nil
	}
	return err
}

// applyRestoredCondition uses server-side apply against the status subresource.
// Pod conditions are an associative list keyed by type, so this field manager
// owns only nvidia.com/Restored and does not replace kubelet-owned conditions.
func (w *NodeController) applyRestoredCondition(ctx context.Context, pod *corev1.Pod, status corev1.ConditionStatus, reason, message string) error {
	setPodCondition(&pod.Status, corev1.PodCondition{
		Type:    corev1.PodConditionType(podcontract.RestoredCondition),
		Status:  status,
		Reason:  reason,
		Message: message,
	})
	restored := findRestoredCondition(pod)
	condition := corev1apply.PodCondition().
		WithType(restored.Type).
		WithStatus(restored.Status).
		WithReason(restored.Reason).
		WithMessage(restored.Message).
		WithLastTransitionTime(restored.LastTransitionTime)
	configuration := corev1apply.Pod(pod.Name, pod.Namespace).
		WithStatus(corev1apply.PodStatus().WithConditions(condition))
	_, err := w.clientset.CoreV1().Pods(pod.Namespace).ApplyStatus(ctx, configuration, metav1.ApplyOptions{
		FieldManager: restoreStatusFieldManager,
		Force:        true,
	})
	return err
}

func (w *NodeController) addRestoreFinalizer(ctx context.Context, pod *corev1.Pod) error {
	if hasFinalizer(pod, restorePodFinalizer) {
		return nil
	}
	finalizers := append(append([]string{}, pod.Finalizers...), restorePodFinalizer)
	if err := w.patchRestoreFinalizers(ctx, pod, finalizers); err != nil {
		return fmt.Errorf("add restore protection finalizer: %w", err)
	}
	pod.Finalizers = finalizers
	return nil
}

func (w *NodeController) removeRestoreFinalizer(ctx context.Context, pod *corev1.Pod) error {
	if !hasFinalizer(pod, restorePodFinalizer) {
		return nil
	}
	finalizers := make([]string, 0, len(pod.Finalizers)-1)
	for _, finalizer := range pod.Finalizers {
		if finalizer != restorePodFinalizer {
			finalizers = append(finalizers, finalizer)
		}
	}
	if err := w.patchRestoreFinalizers(ctx, pod, finalizers); err != nil {
		if apierrors.IsNotFound(err) {
			return nil
		}
		return fmt.Errorf("remove restore protection finalizer: %w", err)
	}
	pod.Finalizers = finalizers
	return nil
}

func (w *NodeController) patchRestoreFinalizers(ctx context.Context, pod *corev1.Pod, finalizers []string) error {
	patch, err := json.Marshal(map[string]any{"metadata": map[string]any{"finalizers": finalizers}})
	if err != nil {
		return err
	}
	_, err = w.clientset.CoreV1().Pods(pod.Namespace).Patch(ctx, pod.Name, ktypes.MergePatchType, patch, metav1.PatchOptions{})
	return err
}

func (w *NodeController) restoreHandled(pod *corev1.Pod) bool {
	_, handled := w.handledRestores.Load(string(pod.UID))
	return handled
}

func (w *NodeController) markRestoreHandled(pod *corev1.Pod) {
	w.handledRestores.Store(string(pod.UID), struct{}{})
}

func (w *NodeController) removeRestoreFinalizerWithEvent(ctx context.Context, pod *corev1.Pod) bool {
	if err := w.removeRestoreFinalizer(ctx, pod); err != nil {
		w.log.Error(err, "Failed to remove restore protection finalizer", "pod", fmt.Sprintf("%s/%s", pod.Namespace, pod.Name))
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreFinalizerUpdateFailedReason, err.Error())
		return true
	}
	return false
}

func (w *NodeController) finishRestore(
	ctx context.Context,
	pod *corev1.Pod,
	status corev1.ConditionStatus,
	reason, message string,
) error {
	if err := w.applyRestoredCondition(ctx, pod, status, reason, message); err != nil {
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, fmt.Sprintf("Failed to record %s restore status: %v", reason, err))
		return err
	}
	w.markRestoreHandled(pod)
	eventType := corev1.EventTypeWarning
	if status == corev1.ConditionTrue {
		eventType = corev1.EventTypeNormal
	}
	emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, eventType, reason, message)
	finalizerErr := w.removeRestoreFinalizer(ctx, pod)
	if finalizerErr != nil {
		w.log.Error(finalizerErr, "Failed to remove restore protection finalizer", "pod", fmt.Sprintf("%s/%s", pod.Namespace, pod.Name))
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreFinalizerUpdateFailedReason, finalizerErr.Error())
	}
	return finalizerErr
}

func (w *NodeController) failRestorePod(ctx context.Context, pod *corev1.Pod, cause error) bool {
	podKey := fmt.Sprintf("%s/%s", pod.Namespace, pod.Name)
	w.log.Error(cause, "Restore request failed", "pod", podKey)
	err := w.finishRestore(
		ctx,
		pod,
		corev1.ConditionFalse,
		podcontract.RestoreReasonFailed,
		cause.Error(),
	)
	return err != nil
}

func (w *NodeController) handleRestorePreflightError(ctx context.Context, pod *corev1.Pod, cause error) bool {
	var incompatible *compat.IncompatibleError
	if errors.As(cause, &incompatible) {
		return w.refuseRestore(ctx, pod, incompatible)
	}

	var pending *restorePendingError
	if !errors.As(cause, &pending) {
		return w.failRestorePod(ctx, pod, cause)
	}

	podKey := fmt.Sprintf("%s/%s", pod.Namespace, pod.Name)
	if restoreInProgress(pod) {
		w.log.V(1).Info("Restore remains in progress while a dependency is pending", "pod", podKey, "reason", pending.reason, "message", pending.message)
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeNormal, pending.reason, pending.message)
		return true
	}
	err := w.applyRestoredCondition(ctx, pod, corev1.ConditionFalse, pending.reason, pending.message)
	if err != nil {
		w.log.Error(err, "Failed to apply pending restore condition", "pod", podKey)
		emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeWarning, restoreStatusUpdateFailedReason, fmt.Sprintf("Failed to record pending restore status: %v", err))
	}
	w.log.V(1).Info("Restore preflight is pending", "pod", podKey, "reason", pending.reason, "message", pending.message)
	emitPodEvent(ctx, w.clientset, w.log, pod, snapshotEventComponent, corev1.EventTypeNormal, pending.reason, pending.message)
	return true
}

func (w *NodeController) enqueueContent(obj interface{}) {
	if name, ok := contentNameFromInformerObj(obj); ok {
		w.captureQueue.Add(name)
	}
}

// enqueueCaptureForSourcePod maps a source-pod event back to every work order naming that pod.
// The index can hold several — contents are named after their PodSnapshot's UID, so each snapshot
// of a pod adds one. Which may dump is decided during reconcile.
func (w *NodeController) enqueueCaptureForSourcePod(obj interface{}) {
	pod, ok := podFromInformerObj(obj)
	if !ok {
		return
	}
	objs, err := w.contentIndexer.ByIndex(podRefIndex, pod.Namespace+"/"+pod.Name)
	if err != nil {
		w.log.Error(err, "Failed to look up PodSnapshotContent by source pod", "pod", fmt.Sprintf("%s/%s", pod.Namespace, pod.Name))
		return
	}
	for _, obj := range objs {
		if content, ok := contentFromInformerObj(obj); ok {
			w.captureQueue.Add(content.Name)
		}
	}
}

// runCaptureQueue gives each work order its own goroutine, for the same reasons as the restore
// queue. Per-key exclusion still comes from the queue itself.
func (w *NodeController) runCaptureQueue(ctx context.Context) {
	for {
		name, shutdown := w.captureQueue.Get()
		if shutdown {
			return
		}
		go w.processCaptureQueueItem(ctx, name)
	}
}

// processCaptureQueueItem holds the key for the whole reconcile, dump included; Adds arriving
// meanwhile collapse into one redelivery after Done.
func (w *NodeController) processCaptureQueueItem(ctx context.Context, name string) {
	defer w.captureQueue.Done(name)
	// A panic here would otherwise kill every other capture and restore on
	// this node.
	defer recoverPanic(w.log, "capture work order", func(err error) {
		w.failCaptureAfterPanic(ctx, name, err)
	}, "content", name)

	if err := w.reconcileCapture(ctx, name); err != nil {
		w.log.Error(err, "Failed to reconcile capture work order", "content", name)
		if ctx.Err() == nil {
			w.captureQueue.AddRateLimited(name)
		}
		return
	}
	w.captureQueue.Forget(name)
}

// podRefIndex is the PodSnapshotContent informer index keyed by source pod ("<namespace>/<name>").
const podRefIndex = "byPodRef"

// podRefIndexFunc indexes a PodSnapshotContent by its source pod ("<snapshotRef.namespace>/<source.podRef.name>").
// It runs against the dynamic informer's *unstructured.Unstructured objects; an unexpected type or a
// missing field yields no index entry (nil) rather than an error, so it never poisons the index.
func podRefIndexFunc(obj interface{}) ([]string, error) {
	u, ok := obj.(*unstructured.Unstructured)
	if !ok {
		return nil, nil
	}
	ns, _, _ := unstructured.NestedString(u.Object, "spec", "snapshotRef", "namespace")
	name, _, _ := unstructured.NestedString(u.Object, "spec", "source", "podRef", "name")
	if ns == "" || name == "" {
		return nil, nil
	}
	return []string{ns + "/" + name}, nil
}

// contentFromInformerObj converts a dynamic informer object (or its DeletedFinalStateUnknown
// tombstone) to a typed PodSnapshotContent. It returns false on an unexpected type.
func contentFromInformerObj(obj interface{}) (*snapshotv1alpha1.PodSnapshotContent, bool) {
	if tombstone, ok := obj.(cache.DeletedFinalStateUnknown); ok {
		obj = tombstone.Obj
	}
	u, ok := obj.(*unstructured.Unstructured)
	if !ok {
		return nil, false
	}
	content := &snapshotv1alpha1.PodSnapshotContent{}
	if err := runtime.DefaultUnstructuredConverter.FromUnstructured(u.Object, content); err != nil {
		return nil, false
	}
	return content, true
}

// chooseActiveContent returns the name of the oldest non-terminal PodSnapshotContent among the indexed
// objects (oldest first by CreationTimestamp, ties broken by Name), or "" when none are active.
// Driving the oldest until it finishes gives deterministic, stable selection across pod events.
func chooseActiveContent(objs []interface{}) string {
	var chosen *snapshotv1alpha1.PodSnapshotContent
	for _, obj := range objs {
		content, ok := contentFromInformerObj(obj)
		if !ok || !content.DeletionTimestamp.IsZero() || isContentTerminal(content) {
			continue
		}
		if chosen == nil ||
			content.CreationTimestamp.Before(&chosen.CreationTimestamp) ||
			(content.CreationTimestamp.Equal(&chosen.CreationTimestamp) && content.Name < chosen.Name) {
			chosen = content
		}
	}
	if chosen == nil {
		return ""
	}
	return chosen.Name
}

func (w *NodeController) restoreArtifactReady(log logr.Logger, podKey, artifactPath string) (bool, error) {
	info, err := os.Stat(artifactPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.V(1).Info("Artifact not ready on disk, skipping restore", "pod", podKey, "artifact_path", artifactPath)
			return false, nil
		}
		return false, fmt.Errorf("stat artifact path %s: %w", artifactPath, err)
	}
	if !info.IsDir() {
		return false, fmt.Errorf("artifact path %s is not a directory", artifactPath)
	}
	return true, nil
}

// restoreContainerIDFromStatus is the fast path when kubelet has already
// published the target container ID. The restore controller polls the node
// runtime when status publication lags container creation.
func restoreContainerIDFromStatus(pod *corev1.Pod, containerName string) string {
	for _, status := range pod.Status.ContainerStatuses {
		if status.Name == containerName && status.ContainerID != "" && status.State.Running != nil {
			return snapshotruntime.StripCRIScheme(status.ContainerID)
		}
	}
	return ""
}

func restoreContainerResolveAttemptContext(ctx context.Context, deadlineAt time.Time) (context.Context, context.CancelFunc) {
	attemptDeadline := time.Now().Add(containerResolveAttemptTimeout)
	if deadlineAt.Before(attemptDeadline) {
		attemptDeadline = deadlineAt
	}
	return context.WithDeadline(ctx, attemptDeadline)
}

func findRestoredCondition(pod *corev1.Pod) *corev1.PodCondition {
	for i := range pod.Status.Conditions {
		condition := &pod.Status.Conditions[i]
		if condition.Type == corev1.PodConditionType(podcontract.RestoredCondition) {
			return condition
		}
	}
	return nil
}

func restoreInProgress(pod *corev1.Pod) bool {
	condition := findRestoredCondition(pod)
	return condition != nil && condition.Status == corev1.ConditionFalse && condition.Reason == podcontract.RestoreReasonInProgress
}

func isRestoreSucceeded(pod *corev1.Pod) bool {
	return podcontract.ClassifyRestoreOutcome(pod.Status.Conditions) == podcontract.RestoreOutcomeSucceeded
}

// isRestorePartiallySucceeded reports the terminal mixed worker outcome.
func isRestorePartiallySucceeded(pod *corev1.Pod) bool {
	return podcontract.ClassifyRestoreOutcome(pod.Status.Conditions) == podcontract.RestoreOutcomePartiallySucceeded
}

func isRestoreTerminal(pod *corev1.Pod) bool {
	return podcontract.ClassifyRestoreOutcome(pod.Status.Conditions).Terminal()
}

func isRestorePodActive(pod *corev1.Pod) bool {
	return pod.Status.Phase == corev1.PodPending || pod.Status.Phase == corev1.PodRunning
}

func hasFinalizer(pod *corev1.Pod, finalizer string) bool {
	for _, item := range pod.Finalizers {
		if item == finalizer {
			return true
		}
	}
	return false
}

func newRestorePendingError(reason, message string) error {
	return &restorePendingError{reason: reason, message: message}
}