import type { Accelerator, DeviceId, LaunchPlan } from "../contracts";

/**
 * The single place device selection becomes a mechanism.
 *
 * Engines declare `plan.devices` abstractly; this translates that list into whatever the
 * accelerator and runtime actually need. Before this existed the same derivation lived in
 * four modules with four different key precedences (gpu-leases, process-utilities,
 * process-manager, backend-builder) — one of which resolved UUIDs against an empty GPU
 * list and so only ever saw the raw selector.
 */

/** Docker flags a launcher must add for this plan. Empty for process launches. */
export interface DeviceRuntimeFlags {
  readonly args: readonly string[];
  readonly groupAdd: readonly string[];
}

const joined = (devices: readonly DeviceId[]): string => devices.join(",");

/** Indices for accelerators whose tooling selects by ordinal rather than UUID. */
const ordinals = (devices: readonly DeviceId[]): string =>
  devices.map((device) => device.slice(device.lastIndexOf(":") + 1)).join(",");

export const deviceEnvironment = (
  accelerator: Accelerator,
  devices: readonly DeviceId[],
): Readonly<Record<string, string>> => {
  if (devices.length === 0) return {};
  switch (accelerator) {
    case "cuda":
      return { CUDA_VISIBLE_DEVICES: joined(devices) };
    case "rocm":
      // ROCR_ gates the runtime, HIP_ gates the HIP API; setting only one leaves the
      // other seeing every card on the box.
      return { HIP_VISIBLE_DEVICES: ordinals(devices), ROCR_VISIBLE_DEVICES: ordinals(devices) };
    case "xpu":
      return { ONEAPI_DEVICE_SELECTOR: `level_zero:${ordinals(devices)}` };
    case "metal":
    case "cpu":
      // Metal exposes no device selection, and CPU has nothing to select.
      return {};
  }
};

export const dockerFlagsFor = (
  accelerator: Accelerator,
  devices: readonly DeviceId[],
): DeviceRuntimeFlags => {
  if (devices.length === 0) return { args: [], groupAdd: [] };
  switch (accelerator) {
    case "cuda":
      return { args: ["--gpus", `"device=${joined(devices)}"`], groupAdd: [] };
    case "rocm":
      return {
        args: ["--device", "/dev/kfd", "--device", "/dev/dri", "--security-opt", "seccomp=unconfined"],
        groupAdd: ["video", "render"],
      };
    case "xpu":
      return { args: ["--device", "/dev/dri"], groupAdd: ["render"] };
    case "metal":
    case "cpu":
      return { args: [], groupAdd: [] };
  }
};

/**
 * Fold device selection into a plan's environment. Runtime is taken from the plan: a
 * docker plan still gets the env vars (the container reads them) *plus* the flags the
 * launcher applies separately via `dockerFlagsFor`.
 */
export const applyDevices = (plan: LaunchPlan, accelerator: Accelerator): LaunchPlan => ({
  ...plan,
  env: { ...plan.env, ...deviceEnvironment(accelerator, plan.devices) },
});