import type {
  EngineSupport,
  HealthCheck,
  LaunchPlan,
  LaunchRequest,
  MetricMap,
  EngineRuntimeKind,
  ServingOptions,
} from "../contracts";

/** Model directory inside a container. Every image mounts the model at the same path, so
 *  the argv is identical whether a plan runs as a process or a container. */
export const CONTAINER_MODEL_DIR = "/models";

export const health = (path: string, readyDeadlineMs: number, intervalMs = 2_000): HealthCheck => ({
  path,
  readyDeadlineMs,
  intervalMs,
});

export const unsupported = (reason: string): EngineSupport => ({ ok: false, reason });
export const supported = (...runtimes: EngineRuntimeKind[]): EngineSupport => ({ ok: true, runtimes });

export const noMetrics: MetricMap = {
  requestsRunning: [],
  requestsWaiting: [],
  kvCacheUtilization: [],
  promptTokensTotal: [],
  generationTokensTotal: [],
};

export const prometheusMetrics = (prefix: string, kvName: string): MetricMap => ({
  requestsRunning: [`${prefix}:num_requests_running`],
  requestsWaiting: [`${prefix}:num_requests_waiting`],
  kvCacheUtilization: [`${prefix}:${kvName}`],
  promptTokensTotal: [`${prefix}:prompt_tokens_total`],
  generationTokensTotal: [`${prefix}:generation_tokens_total`],
});

/* ── tuning knobs ────────────────────────────────────────────────────────── */

/**
 * How one engine spells one canonical knob. `null` in a spelling table means the engine
 * has no equivalent, and the knob is dropped rather than guessed at.
 *
 * This table plus `tuningArguments` is what collapses the two structurally identical 40-line
 * argument builders (vllm-spec.ts:107-148 and sglang-spec.ts:46-92) into data.
 */
export interface FlagSpec {
  readonly flag: string;
  /** Emitted alongside the flag when the knob is set — vLLM's tool parser needs
   *  `--enable-auto-tool-choice` next to it, SGLang's does not. */
  readonly companion?: string;
}

export type TuningKey = keyof ServingOptions;
export type Spelling = Readonly<Partial<Record<TuningKey, FlagSpec>>>;

/** Fixed emission order, so two engines with the same knobs produce comparable argv. */
const TUNING_ORDER: readonly TuningKey[] = [
  "tensorParallel",
  "pipelineParallel",
  "maxContextLength",
  "memoryFraction",
  "maxConcurrentRequests",
  "kvCacheDtype",
  "dtype",
  "quantization",
  "trustRemoteCode",
  "toolCallParser",
  "reasoningParser",
];

/** Parallelism of 1 is the default everywhere; emitting it only adds noise and, for some
 *  builds, forces a distributed code path that a single card does not need. */
const PARALLEL_KEYS = new Set<TuningKey>(["tensorParallel", "pipelineParallel"]);

const shouldEmit = (key: TuningKey, value: ServingOptions[TuningKey]): boolean => {
  if (value === null || value === undefined || value === false) return false;
  if (value === "auto") return false;
  if (typeof value === "number") return PARALLEL_KEYS.has(key) ? value > 1 : value > 0;
  if (typeof value === "string") return value.length > 0;
  return true;
};

export const tuningArguments = (options: ServingOptions, spelling: Spelling): string[] => {
  const args: string[] = [];
  for (const key of TUNING_ORDER) {
    const spec = spelling[key];
    const value = options[key];
    if (!spec || !shouldEmit(key, value)) continue;
    if (typeof value === "boolean") args.push(spec.flag);
    else args.push(spec.flag, String(value));
    if (spec.companion) args.push(spec.companion);
  }
  return args;
};

/** The flag key a token represents, or null when it is a value rather than a flag. */
const flagKey = (token: string): string | null =>
  token.startsWith("--") ? (token.split("=")[0] ?? token).slice(2) : null;

/**
 * Append recipe overrides so they always win: any base flag the user also supplied is
 * dropped first. Without this, both spellings reach the engine and which one applies is
 * left to argparse.
 */
export const mergeArguments = (base: readonly string[], extra: readonly string[]): string[] => {
  const overridden = new Set(
    extra.map(flagKey).filter((key): key is string => key !== null),
  );
  const merged: string[] = [];
  for (let index = 0; index < base.length; index += 1) {
    const token = base[index] ?? "";
    const key = flagKey(token);
    if (key === null || !overridden.has(key)) {
      merged.push(token);
      continue;
    }
    // Skip the flag and its value, if it takes one.
    const next = base[index + 1];
    if (next !== undefined && flagKey(next) === null && !token.includes("=")) index += 1;
  }
  return [...merged, ...extra];
};

/* ── plan assembly ───────────────────────────────────────────────────────── */

export const modelReference = (request: LaunchRequest): string =>
  request.runtime === "docker" ? CONTAINER_MODEL_DIR : request.modelPath;

export const modelMounts = (request: LaunchRequest): LaunchPlan["mounts"] =>
  request.runtime === "docker"
    ? [{ from: request.modelPath, to: CONTAINER_MODEL_DIR, readOnly: true }]
    : [];

/** Containers listen on all interfaces so the published port reaches them; processes bind
 *  loopback, because the controller proxies them and nothing else should connect. */
export const serveAddress = (request: LaunchRequest, listenPort: number): string[] => [
  "--host",
  request.runtime === "docker" ? "0.0.0.0" : "127.0.0.1",
  "--port",
  String(listenPort),
];

/**
 * The shape every OpenAI-compatible server shares. `modelFlag: null` passes the model
 * positionally (vLLM's `serve <path>` form).
 */
export const serverArguments = (
  request: LaunchRequest,
  spec: {
    readonly subcommand?: readonly string[];
    readonly modelFlag: string | null;
    readonly servedNameFlag: string | null;
    readonly spelling: Spelling;
    readonly defaults?: readonly string[];
  },
  listenPort: number,
): string[] => {
  const model = modelReference(request);
  const base = [
    ...(spec.subcommand ?? []),
    ...(spec.modelFlag === null ? [model] : [spec.modelFlag, model]),
    ...(spec.servedNameFlag ? [spec.servedNameFlag, request.servedModelName] : []),
    ...serveAddress(request, listenPort),
    ...tuningArguments(request.options, spec.spelling),
    ...(spec.defaults ?? []),
  ];
  return mergeArguments(base, request.extraArgs);
};

export const plan = (
  request: LaunchRequest,
  parts: {
    readonly args: readonly string[];
    readonly health: HealthCheck;
    readonly listenPort: number;
    readonly image?: string | null;
    readonly env?: Readonly<Record<string, string>>;
  },
): LaunchPlan => {
  const image = request.dockerImage ?? parts.image;
  return {
    kind: request.runtime,
    // A container image supplies its own executable; a process launch needs the binary.
    argv: request.runtime === "docker" ? [...parts.args] : [request.binary, ...parts.args],
    // An engine may always offer an image; only a container plan carries one.
    ...(request.runtime === "docker" && image ? { image } : {}),
    env: { ...request.env, ...(parts.env ?? {}) },
    ports: [{ container: parts.listenPort, host: request.port }],
    mounts: modelMounts(request),
    devices: request.devices,
    health: parts.health,
  };
};