import { AsyncLocalStorage } from "node:async_hooks";
import type { ComposeSpecification } from "../schema/docker.d";

export type ComposeArch = "amd64" | "arm64";

const composeArch = new AsyncLocalStorage<ComposeArch>();

export function resolveComposeArch(compose: ComposeSpecification): ComposeArch {
  return compose["x-arch"] === "amd64" ? "amd64" : "arm64";
}

export function withComposeArch<T>(compose: ComposeSpecification, fn: () => T): T {
  return composeArch.run(resolveComposeArch(compose), fn);
}

export function getComposeArch(): ComposeArch {
  return composeArch.getStore() ?? "arm64";
}

export function getComposeArchPlacement(compose: ComposeSpecification) {
  if (resolveComposeArch(compose) !== "amd64") return {};

  return {
    nodeSelector: {
      "kubernetes.io/arch": "amd64",
    },
    tolerations: [
      {
        key: "arch",
        operator: "Equal",
        value: "amd64",
        effect: "NoExecute",
      },
    ],
  };
}