import { existsSync } from "node:fs";
import { join } from "node:path";
import type { Recipe } from "../../models/types";
import type { Config } from "../../../config/env";
import {
  getUnknownVllmExtraArgKeys as getUnknownVllmExtraArgumentKeys,
  isInternalRecipeKey,
  isJsonStringArgumentKey,
  looksLikeNotesKey,
} from "@local-studio/contracts/engine-args";
import type { Logger } from "../../../core/logger";
import { resolveBinary } from "../../../core/command";
import { managedLlamaServerPath } from "../runtimes/managed-llamacpp";
import { getEngineSpec } from "../engine-spec";
import { resolveRecipeGpuUuids } from "../../system/gpu-leases";

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 const getExtraArgument = (extraArguments: Record<string, unknown>, key: string): unknown => {
  if (Object.prototype.hasOwnProperty.call(extraArguments, key)) {
    return extraArguments[key];
  }
  const kebab = key.replace(/_/g, "-");
  if (Object.prototype.hasOwnProperty.call(extraArguments, kebab)) {
    return extraArguments[kebab];
  }
  const snake = key.replace(/-/g, "_");
  if (Object.prototype.hasOwnProperty.call(extraArguments, snake)) {
    return extraArguments[snake];
  }
  return undefined;
};
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[] => {
  for (const [key, value] of Object.entries(extraArguments)) {
    const normalizedKey = key.replace(/-/g, "_").toLowerCase();
    if (isInternalRecipeKey(key)) {
      continue;
    }
    const flag = `--${key.replace(/_/g, "-")}`;
    if (command.includes(flag)) {
      continue;
    }
    if (value === true) {
      command.push(flag);
      continue;
    }
    if (value === false) {
      if (!["enable_expert_parallelism", "enable-expert-parallelism"].includes(normalizedKey)) {
        command.push(flag);
      }
      continue;
    }
    if (value === undefined || value === null) {
      continue;
    }
    if (typeof value === "string" && isJsonStringArgumentKey(key)) {
      const trimmed = value.trim();
      if (trimmed.startsWith("{") || trimmed.startsWith("[")) {
        try {
          const parsed = JSON.parse(trimmed) as unknown;
          command.push(flag, JSON.stringify(normalizeJsonArgument(parsed)));
          continue;
        } catch {
          command.push(flag, value);
          continue;
        }
      }
    }
    if (Array.isArray(value) || (value && typeof value === "object")) {
      command.push(flag, JSON.stringify(normalizeJsonArgument(value)));
      continue;
    }
    command.push(flag, String(value));
  }
  return command;
};

/**
 * Filter `extraArguments` against the vLLM `serve` flag allowlist and pass the
 * remainder to `appendExtraArguments`. Unknown keys would otherwise be
 * forwarded verbatim, which crashes vLLM with `unrecognized arguments`
 * (real-world example: `benchmark_notes_20260622` blocks the
 * `glm-5-2-504b-term` recipe from booting).
 *
 * Behaviour:
 *   - Unknown keys are dropped unless `LOCAL_STUDIO_ALLOW_UNKNOWN_VLLM_EXTRA_ARGS`
 *     is set to `true` (escape hatch for forked vLLM builds outside the
 *     allowlist).
 *   - Each drop is logged via `logger` (or `console.warn` as a fallback) so the
 *     upstream recipe can be cleaned up.
 *   - Keys that look like free-form notes/annotations are advised to live
 *     under `description` / `metadata` instead.
 */
export const appendVllmExtraArguments = (
  command: string[],
  extraArguments: Record<string, unknown>,
  logger?: Logger,
): string[] => {
  const allowUnknown = process.env["LOCAL_STUDIO_ALLOW_UNKNOWN_VLLM_EXTRA_ARGS"] === "true";
  if (allowUnknown) {
    return appendExtraArguments(command, extraArguments);
  }
  const unknown = getUnknownVllmExtraArgumentKeys(extraArguments);
  if (unknown.length === 0) {
    return appendExtraArguments(command, extraArguments);
  }
  const filtered: Record<string, unknown> = {};
  for (const [key, value] of Object.entries(extraArguments)) {
    if (!unknown.includes(key)) {
      filtered[key] = value;
    }
  }
  const strict = process.env["LOCAL_STUDIO_STRICT_VLLM_EXTRA_ARGS"] === "true";
  for (const key of unknown) {
    const noteLike = looksLikeNotesKey(key);
    const detail: Record<string, unknown> = {
      key,
      hint: noteLike
        ? "vLLM has no such flag; store notes under recipe.description or recipe.metadata"
        : "Add the flag to KNOWN_VLLM_EXTRA_ARG_KEYS in shared/contracts/engine-args.ts, or set LOCAL_STUDIO_ALLOW_UNKNOWN_VLLM_EXTRA_ARGS=true as a temporary escape hatch",
    };
    if (logger) {
      if (strict) {
        logger.error(
          "[vllm-extra-args] dropping unknown vLLM extra_args key in strict mode",
          detail,
        );
      } else {
        logger.warn("[vllm-extra-args] dropping unknown vLLM extra_args key", detail);
      }
    } else if (strict) {
      console.error(
        "[vllm-extra-args] dropping unknown vLLM extra_args key in strict mode",
        detail,
      );
    } else {
      console.warn("[vllm-extra-args] dropping unknown vLLM extra_args key", detail);
    }
  }
  return appendExtraArguments(command, filtered);
};

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;
};

/** In-container path to the vLLM CLI for forked Docker images. */
export const CONTAINER_VLLM_BIN = "/opt/venv/bin/vllm";
const DOCKER_JIT_MOUNT = "/cache/jit";

/**
 * 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 wrapVllmInDocker = (recipe: Recipe, image: string, inner: string[]): string[] => {
  const jitVolume = `local-studio-jit-${sanitizeDockerName(recipe.id)}`;
  return buildDockerRunArguments({
    recipe,
    image,
    inner,
    extraEnv: {
      XDG_CACHE_HOME: DOCKER_JIT_MOUNT,
      CUDA_CACHE_PATH: DOCKER_JIT_MOUNT,
      VLLM_CACHE_DIR: `${DOCKER_JIT_MOUNT}/vllm`,
      TRITON_CACHE_DIR: `${DOCKER_JIT_MOUNT}/triton`,
    },
    extraVolumes: [`${jitVolume}:${DOCKER_JIT_MOUNT}`],
  });
};

const executableBaseName = (value: string): string => {
  return value.split(/[\\/]/).filter(Boolean).at(-1)?.toLowerCase() ?? value.toLowerCase();
};
const isAllowedLlamaServerBinary = (value: string): boolean => {
  const name = executableBaseName(value);
  return name === "llama-server" || name === "llama-server.exe";
};
const rejectPathTraversal = (value: string, label: string): void => {
  if (value.split(/[\\/]+/).includes("..")) {
    throw new Error(`Invalid ${label}: path traversal is not allowed`);
  }
};
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);
};
export const resolveLlamaBinary = (recipe: Recipe, config: Config): string => {
  const override = getExtraArgument(recipe.extra_args, "llama_bin") ?? config.llama_bin;
  if (typeof override === "string" && override.trim()) {
    rejectPathTraversal(override, "llama_bin");
    if (!isAllowedLlamaServerBinary(override)) {
      throw new Error("Invalid llama_bin: only llama-server executables are allowed");
    }
    const resolved = resolveBinary(override);
    if (resolved) {
      return resolved;
    }
    throw new Error(`Invalid llama_bin: executable "${override}" was not found`);
  }
  const managed = managedLlamaServerPath(config);
  return resolveBinary("llama-server") ?? (existsSync(managed) ? managed : "llama-server");
};
export const appendLlamacppArguments = (
  command: string[],
  extraArguments: Record<string, unknown>,
): string[] => {
  for (const [key, value] of Object.entries(extraArguments)) {
    if (isInternalRecipeKey(key)) {
      continue;
    }
    const flag = `--${key.replace(/_/g, "-")}`;
    if (command.includes(flag)) {
      continue;
    }
    if (value === true) {
      command.push(flag);
      continue;
    }
    if (value === false) {
      continue;
    }
    if (value === undefined || value === null || value === "") {
      continue;
    }
    if (Array.isArray(value)) {
      for (const entry of value) {
        if (entry === undefined || entry === null || entry === "") {
          continue;
        }
        command.push(flag, String(entry));
      }
      continue;
    }
    if (typeof value === "object") {
      command.push(flag, JSON.stringify(value));
      continue;
    }
    command.push(flag, String(value));
  }
  return command;
};