#!/usr/bin/env node
// Rebalances the test:e2e:ci-smoke:<n> spec lists in package.json from the
// per-test durations Playwright's list reporter prints in recent successful
// Test runs on main. Usage:
//   node scripts/balance-smoke-shards.mjs [--runs id,id] [--count 6] [--repo owner/name] [--write]
import { execFile } from 'node:child_process';
import { readFileSync, writeFileSync } from 'node:fs';
import { dirname, resolve } from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';
import { promisify } from 'node:util';

const execFileAsync = promisify(execFile);
const ROOT = resolve(dirname(fileURLToPath(import.meta.url)), '..');
const PACKAGE_JSON = resolve(ROOT, 'package.json');
const COMBINED_SCRIPT = 'test:e2e:ci-smoke';
const SHARD_SCRIPT = /^test:e2e:ci-smoke:(\d+)$/;
const SHARD_JOB = /^variant-smoke \((\d+)\)$/;
const SPEC_TOKEN = /^e2e\/\S+\.spec\.[cm]?[jt]sx?$/;

const UNIT_MS = { ms: 1, s: 1000, m: 60_000, h: 3_600_000, d: 86_400_000 };
// One finished test result from the list reporter, after the optional
// Actions timestamp: "  ✓  12 [chromium] › e2e/x.spec.ts:59:3 › title (28.2s)".
// A retried test prints one line per attempt, so failed attempts count too.
const RESULT_LINE =
  /^(?:\d{4}-\d\d-\d\dT\S+Z )?\s+[✓✘]\s+\d+\s+\[[^\]]+\]\s+›\s+(e2e\/\S+?\.spec\.[cm]?[jt]sx?):\d+:\d+\s+›.*\((\d+(?:\.\d+)?)(ms|s|m|h|d)\)\s*$/;

export function stripAnsi(text) {
  return text.replace(/\x1B\[[0-9;?]*[ -/]*[@-~]/g, '');
}

export function parseDurationMs(value, unit) {
  return Math.round(Number(value) * UNIT_MS[unit]);
}

export function specSecondsFromLog(log) {
  const totals = new Map();
  for (const line of stripAnsi(log).split(/\r?\n/)) {
    const match = RESULT_LINE.exec(line);
    if (!match) continue;
    const [, spec, value, unit] = match;
    totals.set(spec, (totals.get(spec) ?? 0) + parseDurationMs(value, unit));
  }
  return totals;
}

export function median(values) {
  if (values.length === 0) throw new Error('median of no values');
  const sorted = [...values].sort((a, b) => a - b);
  const mid = Math.floor(sorted.length / 2);
  return sorted.length % 2 ? sorted[mid] : Math.round((sorted[mid - 1] + sorted[mid]) / 2);
}

// runs: one Map<spec, ms> per run. Specs with no measurement get the median
// of the measured specs' medians, flagged so the output can say so.
export function estimateSpecs(specs, runs) {
  const measured = new Map();
  for (const spec of specs) {
    const samples = runs.filter((run) => run.has(spec)).map((run) => run.get(spec));
    if (samples.length) measured.set(spec, { ms: median(samples), samples, estimated: false });
  }
  const fallback = measured.size ? median([...measured.values()].map((entry) => entry.ms)) : 1000;
  return new Map(
    specs.map((spec) => [spec, measured.get(spec) ?? { ms: fallback, samples: [], estimated: true }]),
  );
}

export function shardLoads(shards, msBySpec) {
  return shards.map((shard) => shard.reduce((sum, spec) => sum + msBySpec.get(spec), 0));
}

// Lexicographic order on loads sorted descending: lowering the largest shard
// wins first, and ties on the max are broken by the next largest.
function compareLoadVectors(a, b) {
  const x = [...a].sort((p, q) => q - p);
  const y = [...b].sort((p, q) => q - p);
  for (let i = 0; i < x.length; i++) if (x[i] !== y[i]) return x[i] - y[i];
  return 0;
}

function longestFirst(specs, msBySpec) {
  return [...specs].sort((a, b) => msBySpec.get(b) - msBySpec.get(a) || a.localeCompare(b));
}

function placeLongestFirst(shards, specs, msBySpec) {
  const loads = shardLoads(shards, msBySpec);
  for (const spec of longestFirst(specs, msBySpec)) {
    const target = loads.indexOf(Math.min(...loads));
    shards[target].push(spec);
    loads[target] += msBySpec.get(spec);
  }
  return shards;
}

// Repeats the single move or pairwise swap that most improves the load
// vector until none does or the largest shard is within goodEnoughMs. Each
// accepted step strictly decreases the vector, so this ends.
function improve(start, msBySpec, goodEnoughMs) {
  let shards = start.map((shard) => [...shard].sort());
  for (;;) {
    let best = null;
    let bestLoads = shardLoads(shards, msBySpec);
    if (Math.max(...bestLoads) <= goodEnoughMs) return shards;
    const consider = (candidate) => {
      const loads = shardLoads(candidate, msBySpec);
      if (compareLoadVectors(loads, bestLoads) < 0) {
        best = candidate;
        bestLoads = loads;
      }
    };
    for (let from = 0; from < shards.length; from++) {
      for (const spec of shards[from]) {
        for (let to = 0; to < shards.length; to++) {
          if (to === from || shards[from].length === 1) continue;
          const next = shards.map((shard) => [...shard]);
          next[from] = next[from].filter((s) => s !== spec);
          next[to] = [...next[to], spec].sort();
          consider(next);
        }
        for (let other = from + 1; other < shards.length; other++) {
          for (const swapWith of shards[other]) {
            const next = shards.map((shard) => [...shard]);
            next[from] = [...next[from].filter((s) => s !== spec), swapWith].sort();
            next[other] = [...next[other].filter((s) => s !== swapWith), spec].sort();
            consider(next);
          }
        }
      }
    }
    if (!best) return shards;
    shards = best;
  }
}

function assertPartition(shards, specs, label) {
  const flat = shards.flat();
  const expected = [...specs].sort();
  const actual = [...flat].sort();
  if (new Set(flat).size !== flat.length || JSON.stringify(actual) !== JSON.stringify(expected)) {
    throw new Error(`${label}: shards do not partition the smoke spec set exactly`);
  }
  if (shards.some((shard) => shard.length === 0)) throw new Error(`${label}: a shard is empty`);
}

// A spec's median still moves several percent between six-run windows, so a
// split within this fraction of the best possible max is as good as optimal.
const GOOD_ENOUGH_FRACTION = 0.01;

// Packs specs into shardCount lists. Two starts are improved: longest
// processing time first, and the current split. The current split wins unless
// the other start lowers the max further than it can, so a drift moves few specs.
export function packShards(specs, msBySpec, shardCount, current = []) {
  if (specs.length < shardCount) throw new Error(`need at least ${shardCount} specs, have ${specs.length}`);
  const inventory = new Set(specs);
  const durations = specs.map((spec) => msBySpec.get(spec));
  const lowerBound = Math.max(durations.reduce((a, b) => a + b, 0) / shardCount, ...durations);
  const goodEnoughMs = lowerBound * (1 + GOOD_ENOUGH_FRACTION);
  const fromLpt = improve(
    placeLongestFirst(Array.from({ length: shardCount }, () => []), specs, msBySpec),
    msBySpec,
    goodEnoughMs,
  );
  const kept = Array.from({ length: shardCount }, (_, i) => (current[i] ?? []).filter((s) => inventory.has(s)));
  const placed = new Set(kept.flat());
  const fromCurrent = improve(
    placeLongestFirst(kept, specs.filter((s) => !placed.has(s)), msBySpec),
    msBySpec,
    goodEnoughMs,
  );
  const maxOf = (shards) => Math.max(...shardLoads(shards, msBySpec));
  const currentIsEnough = maxOf(fromCurrent) <= Math.max(goodEnoughMs, maxOf(fromLpt));
  const shards = (currentIsEnough ? fromCurrent : fromLpt).map((shard) => [...shard].sort());
  assertPartition(shards, specs, 'packShards');
  return shards;
}

export function readSmokeScripts(scripts) {
  const combined = scripts[COMBINED_SCRIPT];
  if (!combined) throw new Error(`package.json has no ${COMBINED_SCRIPT} script`);
  const specs = combined.split(/\s+/).filter((token) => SPEC_TOKEN.test(token));
  const shardNames = Object.keys(scripts)
    .filter((name) => SHARD_SCRIPT.test(name))
    .sort((a, b) => Number(SHARD_SCRIPT.exec(a)[1]) - Number(SHARD_SCRIPT.exec(b)[1]));
  const shards = shardNames.map((name) => scripts[name].split(/\s+/).filter((token) => SPEC_TOKEN.test(token)));
  const unknown = shards.flat().filter((spec) => !specs.includes(spec));
  if (unknown.length) throw new Error(`shard scripts name specs missing from ${COMBINED_SCRIPT}: ${unknown.join(' ')}`);
  return { specs, shardNames, shards };
}

// Rewrites only the shard script values in the package.json text, keeping each
// script's non-spec prefix and suffix, the key order and the file formatting.
export function rewriteShardScripts(packageText, shardNames, shards) {
  const scripts = JSON.parse(packageText).scripts;
  let text = packageText;
  shardNames.forEach((name, i) => {
    const tokens = scripts[name].split(' ');
    const first = tokens.findIndex((token) => SPEC_TOKEN.test(token));
    const last = tokens.findLastIndex((token) => SPEC_TOKEN.test(token));
    if (first < 0) throw new Error(`${name} names no spec paths`);
    if (tokens.slice(first, last + 1).some((token) => !SPEC_TOKEN.test(token))) {
      throw new Error(`${name} interleaves options with spec paths; rewrite it by hand`);
    }
    const value = [...tokens.slice(0, first), ...shards[i], ...tokens.slice(last + 1)].join(' ');
    const before = `${JSON.stringify(name)}: ${JSON.stringify(scripts[name])}`;
    if (text.split(before).length !== 2) throw new Error(`could not locate ${name} in package.json`);
    text = text.replace(before, () => `${JSON.stringify(name)}: ${JSON.stringify(value)}`);
  });
  return text;
}

function parseArgs(argv) {
  const args = { count: 6, repo: 'koala73/worldmonitor', runs: null, write: false };
  for (let i = 0; i < argv.length; i++) {
    const flag = argv[i];
    if (flag === '--write') args.write = true;
    else if (flag === '--runs') args.runs = argv[++i].split(',').map((id) => id.trim()).filter(Boolean);
    else if (flag === '--count') args.count = Number(argv[++i]);
    else if (flag === '--repo') args.repo = argv[++i];
    else throw new Error(`unknown argument ${flag}`);
  }
  if (!Number.isInteger(args.count) || args.count < 1) throw new Error('--count must be a positive integer');
  return args;
}

async function gh(args) {
  const { stdout } = await execFileAsync('gh', args, { maxBuffer: 256 * 1024 * 1024 });
  return stdout;
}

async function smokeJobs(repo, runId, shardCount) {
  const out = await gh([
    'api', '--paginate', `repos/${repo}/actions/runs/${runId}/jobs`,
    '--jq', '.jobs[] | [.id, .name, .conclusion] | @tsv',
  ]);
  const jobs = out.trim().split('\n').filter(Boolean).map((line) => line.split('\t'))
    .filter(([, name]) => SHARD_JOB.test(name))
    .map(([id, name, conclusion]) => ({ id, shard: Number(SHARD_JOB.exec(name)[1]), conclusion }));
  const ok = jobs.filter((job) => job.conclusion === 'success');
  return ok.length === shardCount ? ok : null;
}

async function measureRun(repo, runId, shardCount) {
  const jobs = await smokeJobs(repo, runId, shardCount);
  if (!jobs) return null;
  const run = new Map();
  for (const job of jobs) {
    const log = await gh(['api', '--allow-escape-sequences', `repos/${repo}/actions/jobs/${job.id}/logs`]);
    const perSpec = specSecondsFromLog(log);
    if (perSpec.size === 0) throw new Error(`run ${runId} shard ${job.shard}: no list-reporter result lines parsed`);
    for (const [spec, ms] of perSpec) run.set(spec, (run.get(spec) ?? 0) + ms);
  }
  return run;
}

async function collectRuns(args, shardCount) {
  const candidates = args.runs ?? JSON.parse(await gh([
    'run', 'list', '-R', args.repo, '-w', 'Test', '-b', 'main', '--event', 'push', '--status', 'success',
    '-L', String(args.count * 4), '--json', 'databaseId',
  ])).map((run) => String(run.databaseId));
  const used = [];
  for (const runId of candidates) {
    if (!args.runs && used.length === args.count) break;
    const run = await measureRun(args.repo, runId, shardCount);
    if (run) used.push({ runId, run });
    else if (args.runs) throw new Error(`run ${runId} has no ${shardCount} successful variant-smoke shard jobs`);
    else console.error(`skip run ${runId}: variant-smoke shards did not all run`);
  }
  if (used.length === 0) throw new Error('no usable runs');
  return used;
}

const seconds = (ms) => (ms / 1000).toFixed(1);

function describe(label, shards, msBySpec) {
  const loads = shardLoads(shards, msBySpec);
  const total = loads.reduce((a, b) => a + b, 0);
  console.log(`${label}: ${loads.map(seconds).join(' / ')} test-seconds (max ${seconds(Math.max(...loads))} of ${seconds(total)})`);
}

async function main() {
  const args = parseArgs(process.argv.slice(2));
  const packageText = readFileSync(PACKAGE_JSON, 'utf8');
  const { specs, shardNames, shards: current } = readSmokeScripts(JSON.parse(packageText).scripts);
  const used = await collectRuns(args, shardNames.length);
  const estimates = estimateSpecs(specs, used.map((entry) => entry.run));
  const msBySpec = new Map([...estimates].map(([spec, entry]) => [spec, entry.ms]));

  console.log(`runs (${used.length}): ${used.map((entry) => entry.runId).join(',')}`);
  for (const { runId, run } of used) {
    const observed = current.map((shard) => shard.reduce((sum, spec) => sum + (run.get(spec) ?? 0), 0));
    console.log(`  run ${runId} measured on the current split: ${observed.map(seconds).join(' / ')}`);
  }
  console.log('\nper-spec median test-seconds (min-max over runs):');
  for (const spec of longestFirst(specs, msBySpec)) {
    const { ms, samples, estimated } = estimates.get(spec);
    const spread = samples.length ? `${seconds(Math.min(...samples))}-${seconds(Math.max(...samples))}, n=${samples.length}` : 'unmeasured';
    console.log(`  ${seconds(ms).padStart(7)}  ${spec}  (${spread})${estimated ? '  ESTIMATE: no measurement, using the median spec' : ''}`);
  }

  const next = packShards(specs, msBySpec, shardNames.length, current);
  console.log('');
  describe('before', current, msBySpec);
  describe('after ', next, msBySpec);
  next.forEach((shard, i) => console.log(`  ${shardNames[i]}: ${shard.join(' ')}`));

  if (args.write) {
    writeFileSync(PACKAGE_JSON, rewriteShardScripts(packageText, shardNames, next));
    const written = readSmokeScripts(JSON.parse(readFileSync(PACKAGE_JSON, 'utf8')).scripts);
    assertPartition(written.shards, specs, 'package.json after --write');
    console.log(`\nwrote ${shardNames.join(', ')} to package.json`);
  }
}

if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) {
  main().catch((error) => {
    console.error(error.message);
    process.exit(1);
  });
}