// Running cases several times for one or more versions of a machine, and
// comparing two versions case by case. Shared by the CLI (suite, compare,
// optimize, train) and the SDK.

import type { Case } from "./cases.ts";
import { compareTrials, type Comparison } from "./compare.ts";
import { runCase } from "./converse.ts";
import type { Machine } from "./machine.ts";
import { BudgetExceeded } from "./pricing.ts";
import type { Run } from "./types.ts";
import type { Workspace } from "./workspace.ts";

export type Variant = { machine: Machine; model?: string; revision?: string; version?: string };

/**
 * Run every case `trials` times for each variant, graded, with up to
 * `concurrency` runs at once. Returns runs as [variant][case][trial]. A run
 * that fails is reported and left out, so one error does not lose the rest;
 * reaching the budget stops scheduling and fails the call.
 * `onCase` fires as each case's runs (for every variant) complete.
 */
export async function runMatrix(
  workspace: Workspace,
  variants: Variant[],
  cases: Case[],
  trials: number,
  concurrency: number,
  onCase?: (index: number, runs: Run[][]) => void,
  log: (line: string) => void = (line) => console.error(line),
): Promise<Run[][][]> {
  const out = variants.map(() => cases.map((): Run[] => []));
  const remaining = cases.map(() => trials * variants.length);
  const jobs: Array<{ v: number; c: number }> = [];
  for (let c = 0; c < cases.length; c++) for (let t = 0; t < trials; t++) for (let v = 0; v < variants.length; v++) jobs.push({ v, c });
  let stop: unknown;
  let failures = 0;
  await mapLimit(jobs, concurrency, async ({ v, c }) => {
    if (stop === undefined) {
      const { machine, model, revision, version } = variants[v];
      try {
        out[v][c].push(await runCase(workspace, machine, cases[c].input, { model, revision, version, files: cases[c].files, evaluate: true }));
      } catch (error) {
        if (error instanceof BudgetExceeded || (error instanceof Error && /unknown price/.test(error.message))) stop ??= error;
        else {
          failures++;
          log(`  error in case ${c + 1} (${machine.name}): ${error instanceof Error ? error.message : error}`);
        }
      }
    }
    if (--remaining[c] === 0 && stop === undefined) onCase?.(c, out.map((byCase) => byCase[c]));
  });
  if (stop !== undefined) throw stop;
  if (failures) log(`  ${failures} run${failures === 1 ? "" : "s"} failed and ${failures === 1 ? "is" : "are"} left out of the scores.`);
  return out;
}

/** Compare two versions on the cases where both have at least one successful, graded run. */
export function comparePaired(cases: Case[], runsA: Run[][], runsB: Run[][], log: (line: string) => void = (line) => console.error(line)): Comparison {
  const graded = (runs: Run[]) => runs.filter((r) => r.scorecard);
  const a = runsA.map(graded);
  const b = runsB.map(graded);
  const keep = cases.map((_, i) => a[i].length > 0 && b[i].length > 0);
  const dropped = keep.filter((k) => !k).length;
  if (dropped) log(`  ${dropped} case${dropped === 1 ? "" : "s"} left out of the comparison: a version has no successful run there.`);
  if (dropped === cases.length) throw new Error("No case has successful runs from both versions.");
  const pick = <T,>(xs: T[]) => xs.filter((_, i) => keep[i]);
  return compareTrials(pick(cases).map((c) => c.input), pick(a), pick(b));
}

export async function mapLimit<T, R>(items: T[], limit: number, fn: (item: T) => Promise<R>): Promise<R[]> {
  const out = new Array<R>(items.length);
  let next = 0;
  await Promise.all(
    Array.from({ length: Math.min(limit, items.length) }, async () => {
      while (next < items.length) {
        const i = next++;
        out[i] = await fn(items[i]);
      }
    }),
  );
  return out;
}