import { existsSync } from "node:fs";
import { join } from "node:path";
import type { Recipe } from "../../models/types";
import type { Config } from "../../../config/env";
import {
  isInternalRecipeKey,
  isJsonStringArgumentKey,
} from "@local-studio/contracts/engine-args";
import { getEngineSpec } from "../engine-spec";
import { resolveRecipeGpuUuids } from "../../system/gpu-leases";
import { getExtraArgument } from "../argument-utilities";

export { getExtraArgument };

export const normalizeJsonArgument = (value: unknown): unknown => {
  if (Array.isArray(value)) {
    return value.map((item) => normalizeJsonArgument(item));
  }
  if (value && typeof value === "object") {
    const record = value as Record<string, unknown>;
    return Object.fromEntries(
      Object.entries(record).map(([key, entry]) => [
        key.replace(/-/g, "_"),
        normalizeJsonArgument(entry),
      ]),
    );
  }
  return value;
};

export type ExtraArgumentSerializer = (flag: string, key: string, value: unknown) => string[];

export const appendSerializedArguments = (
  command: string[],
  extraArguments: Record<string, unknown>,
  serialize: ExtraArgumentSerializer,
): string[] => {
  for (const [key, value] of Object.entries(extraArguments)) {
    if (isInternalRecipeKey(key)) continue;
    const flag = `--${key.replace(/_/g, "-")}`;
    if (command.includes(flag)) continue;
    command.push(...serialize(flag, key, value));
  }
  return command;
};

const serializeExtraArgument: ExtraArgumentSerializer = (flag, key, value) => {
  if (value === true) return [flag];
  if (value === false) {
    return key.replace(/-/g, "_").toLowerCase() === "enable_expert_parallelism" ? [] : [flag];
  }
  if (value === undefined || value === null) return [];
  if (typeof value === "string" && isJsonStringArgumentKey(key)) {
    const trimmed = value.trim();
    if (trimmed.startsWith("{") || trimmed.startsWith("[")) {
      try {
        return [flag, JSON.stringify(normalizeJsonArgument(JSON.parse(trimmed) as unknown))];
      } catch {
        return [flag, value];
      }
    }
  }
  if (Array.isArray(value) || (value && typeof value === "object")) {
    return [flag, JSON.stringify(normalizeJsonArgument(value))];
  }
  return [flag, String(value)];
};

export const getPythonPath = (recipe: Recipe): string | undefined => {
  if (recipe.python_path && existsSync(recipe.python_path)) {
    return recipe.python_path;
  }
  const venvPath = getExtraArgument(recipe.extra_args, "venv_path");
  if (typeof venvPath === "string") {
    const pythonBin = join(venvPath, "bin", "python");
    if (existsSync(pythonBin)) {
      return pythonBin;
    }
  }
  return undefined;
};
export const appendExtraArguments = (
  command: string[],
  extraArguments: Record<string, unknown>,
): string[] => appendSerializedArguments(command, extraArguments, serializeExtraArgument);

const normalizeLaunchCommand = (command: string): string => {
  return command
    .replace(/\\\s*\n\s*\+?\s*/g, " ")
    .replace(/^\s*\+\s*/gm, "")
    .trim();
};
const splitLaunchCommand = (command: string): string[] => {
  const normalized = normalizeLaunchCommand(command);
  const result: string[] = [];
  let current = "";
  let quote: "'" | '"' | null = null;
  let escaping = false;
  for (const character of normalized) {
    if (escaping) {
      current += character;
      escaping = false;
      continue;
    }
    if (character === "\\") {
      escaping = true;
      continue;
    }
    if (quote) {
      if (character === quote) {
        quote = null;
      } else {
        current += character;
      }
      continue;
    }
    if (character === "'" || character === '"') {
      quote = character;
      continue;
    }
    if (/\s/.test(character)) {
      if (current) {
        result.push(current);
        current = "";
      }
      continue;
    }
    current += character;
  }
  if (escaping) {
    current += "\\";
  }
  if (current) {
    result.push(current);
  }
  return result;
};
const getLaunchCommandOverride = (recipe: Recipe): string[] | null => {
  const override =
    getExtraArgument(recipe.extra_args, "launch_command") ??
    getExtraArgument(recipe.extra_args, "custom_command");
  if (typeof override !== "string" || !override.trim()) {
    return null;
  }
  // A recipe launch_command/custom_command is arbitrary-binary execution as the
  // controller user. Honour it only when the operator has opted in; otherwise
  // ignore the override and build the command from the structured recipe fields.
  if (process.env["LOCAL_STUDIO_ALLOW_CUSTOM_LAUNCH_COMMAND"] !== "true") {
    return null;
  }
  const command = splitLaunchCommand(override);
  return command.length > 0 ? command : null;
};


/**
 * Env keys that must NOT be forwarded into the container; the image's own baked
 * value (sometimes intentionally empty) is required.
 *
 * NOTE: `NCCL_GRAPH_FILE` is deliberately NOT skipped. The voipmonitor "noxml"
 * NCCL build treats an empty `NCCL_GRAPH_FILE` as a fatal error, so recipes set
 * it to `/dev/null` and that override must reach the container.
 */
const DOCKER_ENV_SKIP_KEYS = new Set([
  "CUDA_VISIBLE_DEVICES",
  "NCCL_GRAPH_DUMP_FILE",
  "VLLM_B12X_MLA_EXTEND_MAX_CHUNKS",
]);

export const sanitizeDockerName = (value: string): string => {
  const cleaned = value.replace(/[^a-zA-Z0-9_.-]/g, "-").replace(/^[^a-zA-Z0-9]+/, "");
  return cleaned.length > 0 ? cleaned : "recipe";
};

const buildDockerEnvironmentFlags = (recipe: Recipe): string[] => {
  const flags: string[] = [];
  const seen = new Set<string>();
  const addEnvironment = (source: unknown): void => {
    if (!source || typeof source !== "object") {
      return;
    }
    for (const [key, value] of Object.entries(source as Record<string, unknown>)) {
      if (value === undefined || value === null) continue;
      if (seen.has(key) || DOCKER_ENV_SKIP_KEYS.has(key)) continue;
      seen.add(key);
      flags.push("-e", `${key}=${String(value)}`);
    }
  };
  addEnvironment(recipe.env_vars);
  addEnvironment(getExtraArgument(recipe.extra_args, "env_vars"));
  return flags;
};

export const buildDockerGpuFlags = (recipe: Recipe): string[] => {
  const resolution = resolveRecipeGpuUuids(recipe, []);
  const selector = resolution.selector?.trim() || "";
  if (resolution.source === "recipe" && !selector) return [];
  const request = selector.includes(",") ? `"device=${selector}"` : `device=${selector}`;
  return selector
    ? ["--gpus", request, "-e", `CUDA_VISIBLE_DEVICES=${selector}`]
    : ["--gpus", "all"];
};

export interface DockerRunOptions {
  recipe: Recipe;
  image: string;
  /** The command to run inside the container, after the image reference. */
  inner: string[];
  /** Overrides the derived `local-studio-{recipe.id}` container name — needed
   * whenever more than one container can exist for the same recipe (e.g. an
   * environment, which is keyed by its own id, not the recipe's). */
  containerName?: string;
  /** Extra `-e KEY=VALUE` pairs to set unconditionally (e.g. engine cache dirs). */
  extraEnv?: Record<string, string>;
  /** Extra `-v` volume mounts beyond the model path, each as `source:target[:mode]`. */
  extraVolumes?: string[];
}

/**
 * Shared `docker run` invocation shape for every engine's Docker-backed launch
 * path: foreground container (so the process-manager stop path's SIGTERM/
 * `--rm` teardown applies unchanged), host networking so the engine binds the
 * recipe's port directly, and the model path bind-mounted read-only.
 */
export const buildDockerRunArguments = ({
  recipe,
  image,
  inner,
  containerName,
  extraEnv: extraEnvironment = {},
  extraVolumes = [],
}: DockerRunOptions): string[] => {
  const name = containerName ?? `local-studio-${sanitizeDockerName(recipe.id)}`;
  const model = recipe.model_path;
  const flags = [
    "docker",
    "run",
    "--rm",
    "--name",
    name,
    ...buildDockerGpuFlags(recipe),
    "--network",
    "host",
    "--ipc",
    "host",
    "--shm-size",
    "32g",
    "--ulimit",
    "memlock=-1",
    "--ulimit",
    "stack=67108864",
  ];
  flags.push(...buildDockerEnvironmentFlags(recipe));
  for (const [key, value] of Object.entries(extraEnvironment)) {
    flags.push("-e", `${key}=${value}`);
  }
  flags.push("-v", `${model}:${model}:ro`);
  for (const volume of extraVolumes) {
    flags.push("-v", volume);
  }
  flags.push(image);
  flags.push(...inner);
  return flags;
};

export const buildBackendCommand = (
  recipe: Recipe,
  config: Config,
  managedGpuSelection = false,
): string[] => {
  const launchCommand = getLaunchCommandOverride(recipe);
  if (launchCommand) {
    if (managedGpuSelection) {
      throw new Error("Custom launch commands cannot use managed GPU selection");
    }
    return launchCommand;
  }
  return getEngineSpec(recipe.backend).buildCommand(recipe, config);
};