package runtime

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"syscall"
	"time"

	types "github.com/beam-cloud/beta9/pkg/types"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/rs/zerolog/log"
)

const (
	runscRestoreStateTimeout      = 30 * time.Second
	runscRestoreStatePollInterval = 25 * time.Millisecond
	runscDeleteTimeout            = 5 * time.Second
	runscOutputDrainTimeout       = 5 * time.Second
	runscGPUAnnotation            = "com.beam.gvisor.nvproxy"
	// RunscTotalMemoryAnnotation tells runsc (beam-cloud/gvisor
	// release-20260714.0-beam.9 and later) the total memory, in bytes, to report
	// to the application instead of the sandbox cgroup's limit.
	RunscTotalMemoryAnnotation    = "dev.gvisor.internal.total-memory"
	runscAllowUnsupportedDriver   = "--nvproxy-allow-unsupported-driver"
	cudaCheckpointContainerPath   = "/usr/local/bin/cuda-checkpoint"
	runscCheckpointImageName      = "checkpoint.img"
	runscCheckpointSpecsKey       = "container_specs"
	runscCheckpointMaxMetadata    = 16 * 1024 * 1024
	sandboxCgroupMountDestination = "/sys/fs/cgroup"
)

// runscStateFileMagic opens every runsc state file ("gVisorSF").
var runscStateFileMagic = []byte{0x67, 0x56, 0x69, 0x73, 0x6f, 0x72, 0x53, 0x46}

// Runsc implements Runtime using the gVisor runsc runtime.
type Runsc struct {
	cfg                   Config
	dockerPacketWriteFlag string
	restoreOutput         sync.Map // container ID -> *runscRestoreOutput
}

// runscBundle owns one operation's OCI spec; GPU state is never shared between containers.
type runscBundle struct {
	path string
	spec *specs.Spec
}

type runscState struct {
	ID     string `json:"id"`
	Pid    int    `json:"pid"`
	Status string `json:"status"`
	Bundle string `json:"bundle"`
}

type runscCommandResult struct {
	exitCode int
	err      error
}

type runscRestoreOutput struct {
	process *os.Process
	done    chan struct{}
}

// NewRunsc creates a new runsc (gVisor) runtime
func NewRunsc(cfg Config) (*Runsc, error) {
	if cfg.RunscPath == "" {
		cfg.RunscPath = "runsc"
	}
	if cfg.RunscRoot == "" {
		cfg.RunscRoot = "/run/gvisor"
	}
	if !hasRunscFlag(cfg.RunscExtraArgs, runscAllowUnsupportedDriver) {
		cfg.RunscExtraArgs = append(append([]string(nil), cfg.RunscExtraArgs...), runscAllowUnsupportedDriver)
	}

	// Check if runsc is available
	if _, err := exec.LookPath(cfg.RunscPath); err != nil {
		return nil, ErrRuntimeNotAvailable{
			Runtime: types.ContainerRuntimeGvisor.String(),
			Reason:  "runsc binary not found in PATH",
		}
	}

	return &Runsc{
		cfg:                   cfg,
		dockerPacketWriteFlag: selectDockerPacketWriteFlag(runscFlags(cfg.RunscPath)),
	}, nil
}

func hasRunscFlag(args []string, name string) bool {
	for _, arg := range args {
		if arg == name || strings.HasPrefix(arg, name+"=") {
			return true
		}
	}
	return false
}

func (r *Runsc) Name() string {
	return types.ContainerRuntimeGvisor.String()
}

func (r *Runsc) Capabilities() Capabilities {
	return Capabilities{
		CheckpointRestore: true,
		GPU:               true,
		OOMEvents:         true,
		JoinExistingNetNS: true,
		CDI:               true,
	}
}

// Prepare mutates the OCI spec to be compatible with gVisor
func (r *Runsc) Prepare(ctx context.Context, spec *specs.Spec) error {
	if spec == nil || spec.Linux == nil {
		return fmt.Errorf("spec is nil")
	}

	spec.Linux.Seccomp = nil
	if !hasGPUDevices(spec) {
		delete(spec.Annotations, runscGPUAnnotation)
		spec.Linux.Devices = nil
		return nil
	}
	if spec.Annotations == nil {
		spec.Annotations = make(map[string]string)
	}
	spec.Annotations[runscGPUAnnotation] = "true"
	r.mountCudaCheckpoint(spec)

	if spec.Process == nil || !slices.ContainsFunc(spec.Linux.Devices, func(device specs.LinuxDevice) bool {
		return device.Path == "/dev/nvidiactl"
	}) {
		spec.Linux.Devices = nil
		return nil
	}

	// CDI supplies the devices, driver mounts and linker hooks.
	spec.Process.Env = nvidiaVisibleDevices(spec.Process.Env, "void")
	// nvproxy virtualizes NVIDIA devices; DRM entries are host-specific.
	spec.Linux.Devices = slices.DeleteFunc(spec.Linux.Devices, func(device specs.LinuxDevice) bool {
		return !strings.HasPrefix(device.Path, "/dev/nvidia")
	})
	if spec.Hooks != nil {
		for i, hook := range spec.Hooks.CreateContainer {
			if strings.HasSuffix(hook.Path, "/nvidia-cdi-hook") && !slices.Contains(hook.Env, "GOMAXPROCS=1") {
				spec.Hooks.CreateContainer[i].Env = append(hook.Env, "GOMAXPROCS=1")
			}
		}
	}

	return nil
}

func nvidiaVisibleDevices(env []string, devices string) []string {
	return append(slices.DeleteFunc(env, func(entry string) bool {
		return strings.HasPrefix(entry, types.NvidiaVisibleDevicesEnv+"=")
	}), types.NvidiaVisibleDevicesEnv+"="+devices)
}

// mountCudaCheckpoint bind-mounts cuda-checkpoint binary into the container
func (r *Runsc) mountCudaCheckpoint(spec *specs.Spec) {
	cudaCheckpointPath, err := exec.LookPath("cuda-checkpoint")
	if err != nil {
		return // Not found, skip
	}

	spec.Mounts = append(spec.Mounts, specs.Mount{
		Destination: cudaCheckpointContainerPath,
		Type:        "bind",
		Source:      cudaCheckpointPath,
		Options:     []string{"bind", "ro"},
	})
}

func hasGPUDevices(spec *specs.Spec) bool {
	if spec == nil || spec.Linux == nil {
		return false
	}

	for _, device := range spec.Linux.Devices {
		if strings.HasPrefix(device.Path, "/dev/nvidia") {
			return true
		}
	}

	for key := range spec.Annotations {
		if strings.HasPrefix(key, "cdi.k8s.io") {
			return true
		}
	}

	return false
}

func (r *Runsc) Run(ctx context.Context, containerID, bundlePath string, opts *RunOpts) (int, error) {
	dockerEnabled := opts != nil && opts.DockerEnabled
	bundle, err := loadRunscBundle(bundlePath)
	if err != nil {
		return -1, err
	}

	args := r.baseArgs(dockerEnabled)
	args = append(args, "--final-metrics-log", r.finalMetricsPath(containerID))
	defer os.Remove(r.finalMetricsPath(containerID))
	if bundle.usesGPU() {
		args = append(args, "--nvproxy=true")
	}
	// Attached `run` deletes its cgroup before returning, losing OOM evidence.
	args = append(args, "run", "--detach", "--bundle", bundlePath, containerID)

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} // Kill entire process tree

	// Detached children retain these descriptors. Drain output independently of
	// cmd.Wait, which must return when startup completes, not when output ends.
	var output, stderr OutputWriter
	if opts != nil {
		output, stderr = opts.OutputWriter, opts.ErrorWriter
	}
	stdoutPipe, closeStdout, err := runscOutputPipe(output)
	if err != nil {
		return -1, err
	}
	defer closeStdout()
	if stdoutPipe != nil {
		cmd.Stdout, cmd.Stderr = stdoutPipe, stdoutPipe
	}
	var stderrPipe *os.File
	if stderr != nil {
		var closeStderr func()
		stderrPipe, closeStderr, err = runscOutputPipe(stderr)
		if err != nil {
			return -1, err
		}
		defer closeStderr()
		cmd.Stderr = stderrPipe
	}
	// Delete before draining output, including descriptors held by children.
	defer r.forceDelete(containerID, dockerEnabled)

	cgroupPath := containerID
	if bundle.spec.Linux != nil && bundle.spec.Linux.CgroupsPath != "" {
		cgroupPath = strings.TrimPrefix(bundle.spec.Linux.CgroupsPath, "/")
	}
	counter := &CgroupOOMWatcher{cgroupRoot: cgroupV2Root, cgroupPath: cgroupPath}
	initialOOMs, _ := counter.readOOMKillCount()

	if err := cmd.Start(); err != nil {
		return -1, err
	}

	// Only the child owns the pipe writers after startup.
	if stdoutPipe != nil {
		_ = stdoutPipe.Close()
	}
	if stderrPipe != nil {
		_ = stderrPipe.Close()
	}
	if err := cmd.Wait(); err != nil {
		return -1, err
	}
	// Detached startup has exited; observers must follow the live sentry PID.
	state, err := r.loadState(ctx, containerID)
	if err != nil {
		return -1, err
	}
	if opts != nil && opts.Started != nil && state.Pid > 0 {
		opts.Started <- state.Pid
	}
	return r.waitForExit(ctx, containerID, dockerEnabled, counter, initialOOMs)
}

func (r *Runsc) Exec(ctx context.Context, containerID string, proc specs.Process, opts *ExecOpts) error {
	// Create a temporary process spec file
	procFile, err := os.CreateTemp("", "runsc-process-*.json")
	if err != nil {
		return fmt.Errorf("failed to create process spec: %w", err)
	}
	defer os.Remove(procFile.Name())

	procJSON, err := json.Marshal(proc)
	if err != nil {
		return fmt.Errorf("failed to marshal process spec: %w", err)
	}

	if _, err := procFile.Write(procJSON); err != nil {
		return fmt.Errorf("failed to write process spec: %w", err)
	}
	procFile.Close()

	args := r.baseArgs(false)
	args = append(args, "exec")
	args = append(args, "--process", procFile.Name())
	args = append(args, containerID)

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)

	if opts != nil && opts.OutputWriter != nil {
		cmd.Stdout = opts.OutputWriter
		cmd.Stderr = opts.OutputWriter
	}

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to exec in container: %w", err)
	}

	// Notify that process has started
	if opts != nil && opts.Started != nil {
		select {
		case opts.Started <- cmd.Process.Pid:
		default:
		}
	}

	return cmd.Wait()
}

func (r *Runsc) UpdateResources(ctx context.Context, containerID string, resources *specs.LinuxResources) error {
	if resources == nil {
		return fmt.Errorf("resources cannot be nil")
	}

	file, err := os.CreateTemp("", "runsc-resources-*.json")
	if err != nil {
		return err
	}
	defer os.Remove(file.Name())

	if err := json.NewEncoder(file).Encode(resources); err != nil {
		file.Close()
		return err
	}
	if err := file.Close(); err != nil {
		return err
	}

	args := append(r.baseArgs(false), "update", "--resources", file.Name(), containerID)
	if output, err := exec.CommandContext(ctx, r.cfg.RunscPath, args...).CombinedOutput(); err != nil {
		return fmt.Errorf("failed to update container resources: %w: %s", err, strings.TrimSpace(string(output)))
	}
	return nil
}

func (r *Runsc) Kill(ctx context.Context, containerID string, sig syscall.Signal, opts *KillOpts) error {
	args := r.baseArgs(false)
	args = append(args, "kill")

	if opts != nil && opts.All {
		args = append(args, "--all")
	}

	args = append(args, containerID, fmt.Sprintf("%d", sig))

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)
	return cmd.Run()
}

func (r *Runsc) Delete(ctx context.Context, containerID string, opts *DeleteOpts) error {
	defer os.Remove(r.finalMetricsPath(containerID))
	args := r.baseArgs(false)
	args = append(args, "delete")

	if opts != nil && opts.Force {
		args = append(args, "--force")
	}

	args = append(args, containerID)

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)
	err := cmd.Run()
	if err == nil {
		r.drainRestoreOutput(containerID)
	}
	return err
}

func (r *Runsc) State(ctx context.Context, containerID string) (State, error) {
	state, err := r.loadState(ctx, containerID)
	if err != nil {
		return State{}, err
	}

	return State{
		ID:     state.ID,
		Pid:    state.Pid,
		Status: state.Status,
	}, nil
}

func (r *Runsc) loadState(ctx context.Context, containerID string) (runscState, error) {
	args := r.baseArgs(false)
	args = append(args, "state", containerID)

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			// Newer runsc releases return 128 when the container bundle has
			// already disappeared; older releases returned 1 for this case.
			if exitErr.ExitCode() == 1 ||
				(exitErr.ExitCode() == 128 && strings.Contains(strings.ToLower(stderr.String()), "loading container: file does not exist")) {
				return runscState{}, ErrContainerNotFound{ContainerID: containerID}
			}
		}
		if output := strings.TrimSpace(stderr.String()); output != "" {
			return runscState{}, fmt.Errorf("failed to get state: %w (stderr: %s)", err, output)
		}
		return runscState{}, fmt.Errorf("failed to get state: %w", err)
	}

	var state runscState
	if err := json.Unmarshal(stdout.Bytes(), &state); err != nil {
		return runscState{}, fmt.Errorf("failed to parse state output: %w", err)
	}
	return state, nil
}

func (r *Runsc) Checkpoint(ctx context.Context, containerID string, opts *CheckpointOpts) error {
	if opts == nil {
		return fmt.Errorf("checkpoint options cannot be nil")
	}

	bundle, err := r.containerBundle(ctx, containerID)
	if err != nil {
		return fmt.Errorf("failed to inspect container bundle: %w", err)
	}

	// Ensure directories exist
	if opts.ImagePath != "" {
		if err := os.MkdirAll(opts.ImagePath, 0755); err != nil {
			return fmt.Errorf("failed to create image path: %w", err)
		}
	}
	if opts.WorkDir != "" {
		if err := os.MkdirAll(opts.WorkDir, 0755); err != nil {
			return fmt.Errorf("failed to create work dir: %w", err)
		}
	}

	args := r.baseArgs(false)
	args = append(args, "checkpoint")
	if bundle.usesGPU() {
		args = append(args, "--cuda-checkpoint-path", cudaCheckpointContainerPath)
	}
	if opts.ImagePath != "" {
		args = append(args, "--image-path", opts.ImagePath)
	}

	// NOTE: gVisor's --work-path flag is marked as "ignored" in runsc
	// but we include it for compatibility
	if opts.WorkDir != "" {
		args = append(args, "--work-path", opts.WorkDir)
	}
	if opts.LeaveRunning {
		args = append(args, "--leave-running")
	}

	// NOTE: gVisor doesn't support CRIU-specific flags like:
	// --allow-open-tcp, --skip-in-flight, --link-remap
	// These are runc/CRIU specific and not available in runsc
	args = append(args, containerID)

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)

	// Capture both stdout and stderr for better error reporting
	var stderr bytes.Buffer
	if opts.OutputWriter != nil {
		cmd.Stdout = opts.OutputWriter
		cmd.Stderr = io.MultiWriter(opts.OutputWriter, &stderr)
	} else {
		cmd.Stderr = &stderr
	}

	if err := cmd.Run(); err != nil {
		stderrStr := stderr.String()
		if stderrStr != "" {
			return fmt.Errorf("checkpoint failed: %w (stderr: %s)", err, stderrStr)
		}
		return fmt.Errorf("checkpoint failed: %w", err)
	}

	return nil
}

func (r *Runsc) Restore(ctx context.Context, containerID string, opts *RestoreOpts) (int, error) {
	if opts == nil {
		return -1, fmt.Errorf("restore options cannot be nil")
	}

	bundle, err := loadRunscBundle(opts.BundlePath)
	if err != nil {
		return -1, err
	}

	if err := bundle.alignRestore(opts.ImagePath); err != nil {
		log.Warn().Err(err).Str("container_id", containerID).Msg("failed to align restore spec with checkpoint")
	}

	// Ensure directories exist
	if opts.WorkDir != "" {
		if err := os.MkdirAll(opts.WorkDir, 0755); err != nil {
			return -1, fmt.Errorf("failed to create work dir: %w", err)
		}
	}

	args := r.baseArgs(false)
	if bundle.usesGPU() {
		args = append(args, "--nvproxy=true")
	}
	args = append(args, "--final-metrics-log", r.finalMetricsPath(containerID))
	args = append(args, "restore", "--background", "--direct")
	if opts.ImagePath != "" {
		args = append(args, "--image-path", opts.ImagePath)
	}
	// Note: gVisor uses --work-path (not --work-dir)
	if opts.WorkDir != "" {
		args = append(args, "--work-path", opts.WorkDir)
	}
	if opts.BundlePath != "" {
		args = append(args, "--bundle", opts.BundlePath)
	}
	args = append(args, containerID)

	cmd := exec.CommandContext(ctx, r.cfg.RunscPath, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	// Restore is attached: this timer starts after the CLI exits, so live logs
	// remain open, but an orphan retaining stdout cannot prevent cleanup.
	cmd.WaitDelay = runscOutputDrainTimeout
	cleanupOnFailure := true
	defer func() {
		if cleanupOnFailure {
			r.forceDelete(containerID, false)
			if cmd.Process != nil {
				_ = cmd.Process.Kill()
			}
		}
	}()

	// Capture stderr for better error reporting
	var stderr bytes.Buffer
	if opts.OutputWriter != nil {
		cmd.Stdout = opts.OutputWriter
		cmd.Stderr = io.MultiWriter(opts.OutputWriter, &stderr)
	} else {
		cmd.Stderr = &stderr
	}

	if err := cmd.Start(); err != nil {
		stderrStr := stderr.String()
		if stderrStr != "" {
			return -1, fmt.Errorf("restore failed to start: %w (stderr: %s)", err, stderrStr)
		}
		return -1, err
	}

	output := &runscRestoreOutput{process: cmd.Process, done: make(chan struct{})}
	r.restoreOutput.Store(containerID, output)
	go func() {
		select {
		case <-ctx.Done():
			r.forceDelete(containerID, false)
			_ = cmd.Process.Kill()
		case <-output.done:
		}
	}()
	restoreDone := make(chan runscCommandResult, 1)
	go func() {
		restoreDone <- runscWaitResult(cmd.Wait(), stderr.String(), "restore")
		close(restoreDone)
		close(output.done)
		r.restoreOutput.CompareAndDelete(containerID, output)
	}()

	pid, err := r.waitForRestoredContainerPID(ctx, containerID, restoreDone)
	if err != nil {
		return -1, err
	}

	if err := r.waitForRestore(ctx, containerID); err != nil {
		return -1, err
	}

	if opts.Started != nil {
		select {
		case opts.Started <- pid:
		case <-ctx.Done():
			return -1, ctx.Err()
		}
	}

	cleanupOnFailure = false
	return 0, nil
}

func (r *Runsc) forceDelete(containerID string, dockerEnabled bool) {
	ctx, cancel := context.WithTimeout(context.Background(), runscDeleteTimeout)
	defer cancel()

	args := append(r.baseArgs(dockerEnabled), "delete", "--force", containerID)
	_ = exec.CommandContext(ctx, r.cfg.RunscPath, args...).Run()
}

func (r *Runsc) waitForRestore(ctx context.Context, containerID string) error {
	args := append(r.baseArgs(false), "wait", "--restore", containerID)
	if output, err := exec.CommandContext(ctx, r.cfg.RunscPath, args...).CombinedOutput(); err != nil {
		return fmt.Errorf("failed waiting for restore completion: %w (output: %s)", err, strings.TrimSpace(string(output)))
	}
	return nil
}

func (r *Runsc) waitForRestoredContainerPID(ctx context.Context, containerID string, restoreDone <-chan runscCommandResult) (int, error) {
	ctx, cancel := context.WithTimeout(ctx, runscRestoreStateTimeout)
	defer cancel()

	ticker := time.NewTicker(runscRestoreStatePollInterval)
	defer ticker.Stop()

	var lastErr error
	var restoreResult *runscCommandResult
	for {
		state, err := r.State(ctx, containerID)
		if err == nil && state.Status == types.RuncContainerStatusRunning && state.Pid > 0 {
			return state.Pid, nil
		}
		if err != nil {
			lastErr = err
		} else {
			lastErr = fmt.Errorf("restored container is not running (status: %s, pid: %d)", state.Status, state.Pid)
		}

		select {
		case <-ctx.Done():
			if restoreResult != nil && restoreResult.err != nil {
				return restoreResult.exitCode, restoreResult.err
			}
			return -1, fmt.Errorf("restore succeeded but restored container state was unavailable: %w", lastErr)
		case result := <-restoreDone:
			restoreResult = &result
			restoreDone = nil
			if result.err != nil {
				return result.exitCode, result.err
			}
		case <-ticker.C:
		}
	}
}

func runscWaitResult(err error, stderr, operation string) runscCommandResult {
	if err == nil {
		return runscCommandResult{exitCode: 0}
	}

	if exitErr, ok := err.(*exec.ExitError); ok {
		if ws, ok := exitErr.Sys().(syscall.WaitStatus); ok {
			exitCode := ws.ExitStatus()
			if stderr != "" {
				return runscCommandResult{
					exitCode: exitCode,
					err:      fmt.Errorf("%s failed with exit code %d (stderr: %s)", operation, exitCode, stderr),
				}
			}
			return runscCommandResult{exitCode: exitCode, err: err}
		}
	}

	if stderr != "" {
		return runscCommandResult{
			exitCode: -1,
			err:      fmt.Errorf("%s failed: %w (stderr: %s)", operation, err, stderr),
		}
	}

	return runscCommandResult{exitCode: -1, err: err}
}

func (r *Runsc) containerBundle(ctx context.Context, containerID string) (*runscBundle, error) {
	state, err := r.loadState(ctx, containerID)
	if err != nil {
		return nil, err
	}
	return loadRunscBundle(state.Bundle)
}

func loadRunscBundle(bundlePath string) (*runscBundle, error) {
	if bundlePath == "" {
		return nil, fmt.Errorf("container bundle path is empty")
	}
	config, err := os.ReadFile(filepath.Join(bundlePath, "config.json"))
	if err != nil {
		return nil, fmt.Errorf("read container bundle: %w", err)
	}
	var spec specs.Spec
	if err := json.Unmarshal(config, &spec); err != nil {
		return nil, fmt.Errorf("decode container bundle: %w", err)
	}
	return &runscBundle{path: bundlePath, spec: &spec}, nil
}

func (b *runscBundle) usesGPU() bool {
	return b.spec.Annotations[runscGPUAnnotation] == "true" || hasGPUDevices(b.spec)
}

func (b *runscBundle) save() error {
	config, err := json.Marshal(b.spec)
	if err != nil {
		return fmt.Errorf("encode restore bundle: %w", err)
	}
	return os.WriteFile(filepath.Join(b.path, "config.json"), config, 0644)
}

// alignRestore preserves runtime-owned mount settings and the saved NVIDIA
// device layout. Other mounts remain subject to runsc's restore validation.
func (b *runscBundle) alignRestore(imagePath string) error {
	if imagePath == "" {
		return nil
	}

	saved, err := readRunscCheckpointSpec(imagePath)
	if err != nil {
		return err
	}
	checkpointMount, checkpointHasMount := findMount(saved.Mounts, sandboxCgroupMountDestination)

	spec := b.spec
	restoreMount, restoreHasMount := findMount(spec.Mounts, sandboxCgroupMountDestination)
	devicesCleared := spec.Linux != nil && len(spec.Linux.Devices) > 0 && spec.Annotations[runscGPUAnnotation] == "true" &&
		saved.Annotations[runscGPUAnnotation] == "true" && (saved.Linux == nil || len(saved.Linux.Devices) == 0)
	if devicesCleared {
		// Older checkpoints saved no NVIDIA device entries.
		spec.Linux.Devices = nil
	}
	mountChanged := checkpointHasMount != restoreHasMount || !mountsEqual(checkpointMount, restoreMount)
	processManagerChanged := false
	if savedManager, ok := findMount(saved.Mounts, types.WorkerSandboxProcessManagerContainerPath); ok {
		for i, mount := range spec.Mounts {
			// The saved binary moves from the worker overlay to checkpoint storage.
			// Preserve the gofer cache option that was saved with its executable mappings.
			if mount.Destination == savedManager.Destination && mount.Source != types.WorkerSandboxProcessManagerWorkerPath &&
				slices.Contains(savedManager.Options, "overlayfs_stale_read") && !slices.Contains(mount.Options, "overlayfs_stale_read") {
				spec.Mounts[i].Options = append(slices.Clone(mount.Options), "overlayfs_stale_read")
				processManagerChanged = true
			}
		}
	}
	if !devicesCleared && !mountChanged && !processManagerChanged {
		return nil
	}
	if devicesCleared && spec.Process != nil {
		// Older Beam workers added only CDI host groups to the empty base group list.
		// Preserve those saved credentials without overriding requested groups or IDs.
		if saved.Process != nil && len(spec.Process.User.AdditionalGids) == 0 {
			spec.Process.User.AdditionalGids = saved.Process.User.AdditionalGids
		}
		for _, env := range spec.Process.Env {
			if devices, ok := strings.CutPrefix(env, types.WorkerGPUDevicesEnv+"="); ok {
				spec.Process.Env = nvidiaVisibleDevices(spec.Process.Env, devices)
				break
			}
		}
	}

	if mountChanged {
		spec.Mounts = slices.DeleteFunc(spec.Mounts, func(m specs.Mount) bool {
			return filepath.Clean(m.Destination) == sandboxCgroupMountDestination
		})
		if checkpointHasMount {
			spec.Mounts = append(spec.Mounts, checkpointMount)
		}
	}

	if err := b.save(); err != nil {
		return err
	}
	log.Info().
		Str("bundle", b.path).
		Bool("checkpoint_mounts_cgroup", checkpointHasMount).
		Bool("devices_cleared", devicesCleared).
		Msg("aligned restore spec with checkpoint")
	return nil
}

// readRunscCheckpointSpec returns the outer container's saved OCI spec.
// A runsc state file opens with an 8-byte magic, an 8-byte
// big-endian metadata length and a JSON string map; the specs are the map's
// "container_specs" entry, keyed by container name.
func readRunscCheckpointSpec(imagePath string) (*specs.Spec, error) {
	image, err := os.Open(filepath.Join(imagePath, runscCheckpointImageName))
	if err != nil {
		return nil, fmt.Errorf("open checkpoint image: %w", err)
	}
	defer image.Close()

	header := make([]byte, len(runscStateFileMagic)+8)
	if _, err := io.ReadFull(image, header); err != nil {
		return nil, fmt.Errorf("read checkpoint image header: %w", err)
	}
	if !bytes.Equal(header[:len(runscStateFileMagic)], runscStateFileMagic) {
		return nil, fmt.Errorf("checkpoint image is not a runsc state file")
	}
	metadataLen := binary.BigEndian.Uint64(header[len(runscStateFileMagic):])
	if metadataLen > runscCheckpointMaxMetadata {
		return nil, fmt.Errorf("checkpoint metadata length %d exceeds %d", metadataLen, runscCheckpointMaxMetadata)
	}

	rawMetadata := make([]byte, metadataLen)
	if _, err := io.ReadFull(image, rawMetadata); err != nil {
		return nil, fmt.Errorf("read checkpoint metadata: %w", err)
	}
	var metadata map[string]string
	if err := json.Unmarshal(rawMetadata, &metadata); err != nil {
		return nil, fmt.Errorf("decode checkpoint metadata: %w", err)
	}
	rawSpecs, ok := metadata[runscCheckpointSpecsKey]
	if !ok {
		return nil, fmt.Errorf("checkpoint metadata has no %s", runscCheckpointSpecsKey)
	}
	var containerSpecs map[string]*specs.Spec
	if err := json.Unmarshal([]byte(rawSpecs), &containerSpecs); err != nil {
		return nil, fmt.Errorf("decode checkpoint container specs: %w", err)
	}
	if len(containerSpecs) != 1 {
		return nil, fmt.Errorf("checkpoint requires one container spec, got %d", len(containerSpecs))
	}
	for _, spec := range containerSpecs {
		if spec != nil {
			return spec, nil
		}
	}
	return nil, fmt.Errorf("checkpoint container spec is nil")
}

func findMount(mounts []specs.Mount, destination string) (specs.Mount, bool) {
	for _, m := range mounts {
		if filepath.Clean(m.Destination) == destination {
			return m, true
		}
	}
	return specs.Mount{}, false
}

// mountsEqual compares mounts the way runsc's restore validation does: by
// destination, type and option set, ignoring source.
func mountsEqual(a, b specs.Mount) bool {
	if filepath.Clean(a.Destination) != filepath.Clean(b.Destination) || a.Type != b.Type || len(a.Options) != len(b.Options) {
		return false
	}
	options := make(map[string]int, len(a.Options))
	for _, opt := range a.Options {
		options[opt]++
	}
	for _, opt := range b.Options {
		if options[opt] == 0 {
			return false
		}
		options[opt]--
	}
	return true
}

func (r *Runsc) Close() error {
	return nil
}

// baseArgs returns the common arguments for all runsc commands
func (r *Runsc) baseArgs(dockerEnabled bool) []string {
	args := []string{
		"--root", r.cfg.RunscRoot,
	}

	if r.cfg.Debug {
		args = append(args, "--debug", "--debug-log", filepath.Join(r.cfg.RunscRoot, "debug.log"))
	}

	if r.cfg.RunscPlatform != "" {
		args = append(args, "--platform", r.cfg.RunscPlatform)
	}

	args = append(args, r.cfg.RunscExtraArgs...)

	// Add flags required for Docker to function properly inside gVisor.
	if dockerEnabled {
		args = append(args, "--net-raw")
		if r.dockerPacketWriteFlag != "" {
			args = append(args, r.dockerPacketWriteFlag)
		}
	}

	return args
}

func runscFlags(runscPath string) string {
	out, err := exec.Command(runscPath, "flags").Output()
	if err != nil {
		return ""
	}
	return string(out)
}

func selectDockerPacketWriteFlag(flags string) string {
	switch {
	case strings.Contains(flags, "-allow-packet-socket-write"):
		return "--allow-packet-socket-write"
	case strings.Contains(flags, "-TESTONLY-allow-packet-endpoint-write"):
		return "--TESTONLY-allow-packet-endpoint-write"
	default:
		return ""
	}
}

// AddDockerInDockerCapabilities adds the capabilities required for running Docker inside gVisor.
func (r *Runsc) AddDockerInDockerCapabilities(spec *specs.Spec) {
	AddDockerInDockerCapabilities(spec)
}

// mergeCapabilities merges two capability lists, avoiding duplicates
func mergeCapabilities(existing []string, toAdd []string) []string {
	capSet := make(map[string]bool)

	// Add existing capabilities to the set
	for _, cap := range existing {
		capSet[cap] = true
	}

	// Add new capabilities to the set
	for _, cap := range toAdd {
		capSet[cap] = true
	}

	// Convert set back to slice
	result := make([]string, 0, len(capSet))
	for cap := range capSet {
		result = append(result, cap)
	}

	return result
}