#!/usr/bin/env node
/**
 * Benchmark Ollama embedding models on the current machine.
 *
 * Emits a readable table by default, or JSON with --json.
 */

import { performance } from "node:perf_hooks";
import { parseArgs } from "node:util";

const DEFAULT_TEXT =
  "Generate a stable semantic embedding for this short bilingual memory sample. " +
  "It should be suitable for retrieval, recall, and near-duplicate detection.";

type BenchmarkOk = {
  model: string;
  ok: true;
  single_wall_seconds: number;
  single_api_seconds: number | null;
  batch_wall_seconds: number;
  batch_api_seconds: number | null;
  batch_size: number;
  vector_dims: number;
  rounds: number;
};

type BenchmarkErr = { model: string; ok: false; error: string };
type BenchmarkResult = BenchmarkOk | BenchmarkErr;

function usageText(): string {
  return [
    "Usage:",
    "  node scripts/bench_ollama_embeddings.ts [options] <model...>",
    "",
    "Options:",
    "  --host <url>           Ollama host (default http://127.0.0.1:11434)",
    "  --rounds <n>           Measured rounds per test after warmup (default 1)",
    "  --batch-size <n>       Inputs per batch test (default 4)",
    "  --timeout <seconds>    Request timeout (default 300)",
    "  --single-text <text>   Input used for the single-input benchmark",
    "  --skip-warmup          Skip warmup request per model",
    "  --json                 Emit JSON results",
  ].join("\n");
}

function parseNumber(raw: unknown, name: string): number {
  const value = Number(raw);
  if (!Number.isFinite(value)) {
    throw new Error(`Invalid --${name}: ${String(raw)}`);
  }
  return value;
}

function median(values: number[]): number {
  if (!values.length) {
    throw new Error("median() requires at least one value");
  }
  const sorted = [...values].sort((a, b) => a - b);
  const mid = Math.floor(sorted.length / 2);
  if (sorted.length % 2 === 0) {
    return (sorted[mid - 1]! + sorted[mid]!) / 2;
  }
  return sorted[mid]!;
}

function medianOrNull(values: Array<number | null>): number | null {
  const usable = values.filter((v): v is number => typeof v === "number" && Number.isFinite(v));
  return usable.length ? median(usable) : null;
}

function buildBatchInputs(text: string, batchSize: number): string[] {
  return Array.from({ length: batchSize }, (_, index) => `${text} [sample ${index + 1}]`);
}

async function fetchWithTimeout(url: string, init: RequestInit, timeoutMs: number): Promise<Response> {
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), timeoutMs);
  try {
    return await fetch(url, { ...init, signal: controller.signal });
  } finally {
    clearTimeout(timer);
  }
}

async function postEmbed(params: {
  host: string;
  model: string;
  inputs: string[];
  timeoutSeconds: number;
}): Promise<{ embeddings: unknown; total_duration?: unknown }> {
  const url = params.host.replace(/\/+$/, "") + "/api/embed";
  const payload = {
    model: params.model,
    input: params.inputs.length > 1 ? params.inputs : params.inputs[0],
  };

  const response = await fetchWithTimeout(
    url,
    {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(payload),
    },
    Math.max(1, params.timeoutSeconds) * 1000,
  );

  const text = await response.text();
  if (!response.ok) {
    throw new Error(`HTTP ${response.status}: ${text}`);
  }

  let data: unknown;
  try {
    data = JSON.parse(text);
  } catch {
    throw new Error(`Unexpected JSON from Ollama: ${text.slice(0, 256)}`);
  }
  return data as { embeddings: unknown; total_duration?: unknown };
}

async function runOne(params: {
  host: string;
  model: string;
  inputs: string[];
  timeoutSeconds: number;
}): Promise<{ wall_seconds: number; api_seconds: number | null; items: number; vector_dims: number }> {
  const started = performance.now();
  const response = await postEmbed(params);
  const elapsed = (performance.now() - started) / 1000;

  const totalDuration = (response as { total_duration?: unknown }).total_duration;
  const api_seconds =
    typeof totalDuration === "number" && Number.isFinite(totalDuration)
      ? totalDuration / 1_000_000_000
      : null;

  const embeddings = (response as { embeddings?: unknown }).embeddings;
  if (!Array.isArray(embeddings)) {
    throw new Error(`Unexpected Ollama response for ${params.model}: ${JSON.stringify(response)}`);
  }
  const first = embeddings[0];
  const vector_dims = Array.isArray(first) ? first.length : 0;

  return {
    wall_seconds: elapsed,
    api_seconds,
    items: embeddings.length,
    vector_dims,
  };
}

async function benchmarkModel(params: {
  host: string;
  model: string;
  singleText: string;
  batchInputs: string[];
  rounds: number;
  timeoutSeconds: number;
  skipWarmup: boolean;
}): Promise<BenchmarkOk> {
  if (!params.skipWarmup) {
    await runOne({ host: params.host, model: params.model, inputs: [params.singleText], timeoutSeconds: params.timeoutSeconds });
  }

  const singleRuns = [];
  const batchRuns = [];
  for (let index = 0; index < params.rounds; index += 1) {
    singleRuns.push(await runOne({ host: params.host, model: params.model, inputs: [params.singleText], timeoutSeconds: params.timeoutSeconds }));
  }
  for (let index = 0; index < params.rounds; index += 1) {
    batchRuns.push(await runOne({ host: params.host, model: params.model, inputs: params.batchInputs, timeoutSeconds: params.timeoutSeconds }));
  }

  return {
    model: params.model,
    ok: true,
    single_wall_seconds: median(singleRuns.map((run) => run.wall_seconds)),
    single_api_seconds: medianOrNull(singleRuns.map((run) => run.api_seconds)),
    batch_wall_seconds: median(batchRuns.map((run) => run.wall_seconds)),
    batch_api_seconds: medianOrNull(batchRuns.map((run) => run.api_seconds)),
    batch_size: params.batchInputs.length,
    vector_dims: singleRuns.at(-1)?.vector_dims ?? 0,
    rounds: params.rounds,
  };
}

function formatSeconds(value: number | null): string {
  return value === null ? "-" : value.toFixed(2);
}

function printTable(results: BenchmarkResult[], batchSize: number): void {
  const headers = [
    "Model",
    "Single(s)",
    `Batch${batchSize}(s)`,
    "API Single(s)",
    `API Batch${batchSize}(s)`,
    "Dims",
    "Status",
  ];

  const rows: string[][] = results.map((result) => {
    if (!result.ok) {
      return [result.model, "-", "-", "-", "-", "-", result.error];
    }
    return [
      result.model,
      formatSeconds(result.single_wall_seconds),
      formatSeconds(result.batch_wall_seconds),
      formatSeconds(result.single_api_seconds),
      formatSeconds(result.batch_api_seconds),
      String(result.vector_dims),
      "ok",
    ];
  });

  const widths = headers.map((header, index) =>
    Math.max(header.length, ...rows.map((row) => (row[index] ?? "").length)),
  );

  const emit = (cols: string[]) => {
    process.stdout.write(
      `${cols.map((value, index) => value.padEnd(widths[index] ?? value.length)).join("  ")}\n`,
    );
  };

  emit(headers);
  emit(widths.map((width) => "-".repeat(width)));
  for (const row of rows) {
    emit(row);
  }

  const successful = results.filter((r): r is BenchmarkOk => r.ok);
  if (successful.length) {
    const ranked = [...successful].sort((a, b) => {
      if (a.batch_wall_seconds !== b.batch_wall_seconds) return a.batch_wall_seconds - b.batch_wall_seconds;
      if (a.single_wall_seconds !== b.single_wall_seconds) return a.single_wall_seconds - b.single_wall_seconds;
      return a.model.localeCompare(b.model);
    });
    process.stdout.write(`\nRecommended by speed: ${ranked[0]!.model}\n`);
  }
}

async function main(): Promise<number> {
  const parsed = parseArgs({
    args: process.argv.slice(2),
    allowPositionals: true,
    options: {
      host: { type: "string", default: "http://127.0.0.1:11434" },
      rounds: { type: "string", default: "1" },
      "batch-size": { type: "string", default: "4" },
      timeout: { type: "string", default: "300" },
      "single-text": { type: "string", default: DEFAULT_TEXT },
      "skip-warmup": { type: "boolean", default: false },
      json: { type: "boolean", default: false },
      help: { type: "boolean", default: false },
    },
  });

  if (parsed.values.help) {
    process.stdout.write(`${usageText()}\n`);
    return 0;
  }

  const models = parsed.positionals.map((value) => String(value).trim()).filter(Boolean);
  if (!models.length) {
    process.stderr.write(`${usageText()}\n`);
    return 1;
  }

  const host = String(parsed.values.host ?? "").trim() || "http://127.0.0.1:11434";
  const rounds = parseNumber(parsed.values.rounds, "rounds");
  const batchSize = parseNumber(parsed.values["batch-size"], "batch-size");
  const timeoutSeconds = parseNumber(parsed.values.timeout, "timeout");
  const singleText = String(parsed.values["single-text"] ?? DEFAULT_TEXT);
  const skipWarmup = Boolean(parsed.values["skip-warmup"]);
  const asJson = Boolean(parsed.values.json);

  if (rounds < 1) {
    throw new Error("--rounds must be at least 1");
  }
  if (batchSize < 1) {
    throw new Error("--batch-size must be at least 1");
  }

  const batchInputs = buildBatchInputs(singleText, batchSize);
  const results: BenchmarkResult[] = [];

  for (const model of models) {
    try {
      results.push(
        await benchmarkModel({
          host,
          model,
          singleText,
          batchInputs,
          rounds,
          timeoutSeconds,
          skipWarmup,
        }),
      );
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      results.push({ model, ok: false, error: message });
    }
  }

  const successful = results.some((result) => result.ok);
  if (asJson) {
    process.stdout.write(`${JSON.stringify({ results }, null, 2)}\n`);
  } else {
    printTable(results, batchSize);
  }

  return successful ? 0 : 1;
}

void main()
  .then((code) => {
    process.exitCode = code;
  })
  .catch((err) => {
    const message = err instanceof Error ? err.message : String(err);
    process.stderr.write(`${message}\n`);
    process.exitCode = 1;
  });