#!/usr/bin/env node

import path from "node:path";
import { readdir, readFile, stat, writeFile } from "node:fs/promises";
import { fileURLToPath } from "node:url";

function round(value, places = 4) {
  const factor = 10 ** places;
  return Math.round(value * factor) / factor;
}

export function calculateStats(values) {
  if (!Array.isArray(values) || values.length === 0) {
    return { mean: 0, stddev: 0, min: 0, max: 0 };
  }
  const mean = values.reduce((sum, value) => sum + value, 0) / values.length;
  let stddev = 0;
  if (values.length > 1) {
    const variance = values.reduce((sum, value) => sum + ((value - mean) ** 2), 0) / (values.length - 1);
    stddev = Math.sqrt(variance);
  }
  return {
    mean: round(mean, 4),
    stddev: round(stddev, 4),
    min: round(Math.min(...values), 4),
    max: round(Math.max(...values), 4),
  };
}

async function pathExists(targetPath) {
  try {
    await stat(targetPath);
    return true;
  } catch {
    return false;
  }
}

async function readJsonIfExists(targetPath) {
  if (!(await pathExists(targetPath))) {
    return null;
  }
  try {
    return JSON.parse(await readFile(targetPath, "utf8"));
  } catch {
    return null;
  }
}

async function listDirectories(targetPath) {
  const entries = await readdir(targetPath, { withFileTypes: true });
  return entries.filter((entry) => entry.isDirectory()).map((entry) => entry.name).sort();
}

export async function loadRunResults(benchmarkDir) {
  const runsDir = path.join(benchmarkDir, "runs");
  const benchmarkEvalDirs = (await pathExists(benchmarkDir))
    ? (await listDirectories(benchmarkDir)).filter((name) => name.startsWith("eval-"))
    : [];
  let searchDir = null;
  if (await pathExists(runsDir)) {
    searchDir = runsDir;
  } else if (benchmarkEvalDirs.length > 0) {
    searchDir = benchmarkDir;
  } else {
    return {};
  }

  const results = {};
  const evalDirs = (await listDirectories(searchDir)).filter((name) => name.startsWith("eval-"));
  for (const [evalIdx, evalDirName] of evalDirs.entries()) {
    const evalDir = path.join(searchDir, evalDirName);
    const metadata = await readJsonIfExists(path.join(evalDir, "eval_metadata.json"));
    const parsedEvalId = Number.parseInt(evalDirName.split("-")[1], 10);
    const evalId = metadata?.eval_id ?? (Number.isNaN(parsedEvalId) ? evalIdx : parsedEvalId);
    for (const configName of await listDirectories(evalDir)) {
      const configDir = path.join(evalDir, configName);
      const childDirs = await listDirectories(configDir);
      const runDirs = childDirs.filter((name) => name.startsWith("run-"));
      if (runDirs.length === 0) {
        continue;
      }
      if (!results[configName]) {
        results[configName] = [];
      }
      for (const runDirName of runDirs) {
        const runDir = path.join(configDir, runDirName);
        const grading = await readJsonIfExists(path.join(runDir, "grading.json"));
        if (!grading) {
          continue;
        }
        const timing = grading.timing ?? {};
        let timeSeconds = Number(timing.total_duration_seconds ?? 0);
        let tokens = 0;
        const timingJson = await readJsonIfExists(path.join(runDir, "timing.json"));
        if (timeSeconds === 0 && timingJson) {
          timeSeconds = Number(timingJson.total_duration_seconds ?? 0);
          tokens = Number(timingJson.total_tokens ?? 0);
        }
        const executionMetrics = grading.execution_metrics ?? {};
        if (tokens === 0) {
          tokens = Number(executionMetrics.output_chars ?? 0);
        }
        results[configName].push({
          eval_id: evalId,
          run_number: Number.parseInt(runDirName.split("-")[1], 10) || 0,
          pass_rate: Number(grading.summary?.pass_rate ?? 0),
          passed: Number(grading.summary?.passed ?? 0),
          failed: Number(grading.summary?.failed ?? 0),
          total: Number(grading.summary?.total ?? 0),
          time_seconds: timeSeconds,
          tokens,
          tool_calls: Number(executionMetrics.total_tool_calls ?? 0),
          errors: Number(executionMetrics.errors_encountered ?? 0),
          expectations: Array.isArray(grading.expectations) ? grading.expectations : [],
          notes: [
            ...(grading.user_notes_summary?.uncertainties ?? []),
            ...(grading.user_notes_summary?.needs_review ?? []),
            ...(grading.user_notes_summary?.workarounds ?? []),
          ],
        });
      }
    }
  }
  return results;
}

export function aggregateResults(results) {
  const runSummary = {};
  const configs = Object.keys(results);
  for (const config of configs) {
    const runs = results[config] ?? [];
    runSummary[config] = {
      pass_rate: calculateStats(runs.map((run) => Number(run.pass_rate ?? 0))),
      time_seconds: calculateStats(runs.map((run) => Number(run.time_seconds ?? 0))),
      tokens: calculateStats(runs.map((run) => Number(run.tokens ?? 0))),
    };
  }
  const primaryConfig = configs[0];
  const baselineConfig = configs[1];
  const primary = runSummary[primaryConfig] ?? {};
  const baseline = runSummary[baselineConfig] ?? {};
  runSummary.delta = {
    pass_rate: `${(primary.pass_rate?.mean ?? 0) - (baseline.pass_rate?.mean ?? 0) >= 0 ? "+" : ""}${round((primary.pass_rate?.mean ?? 0) - (baseline.pass_rate?.mean ?? 0), 2).toFixed(2)}`,
    time_seconds: `${(primary.time_seconds?.mean ?? 0) - (baseline.time_seconds?.mean ?? 0) >= 0 ? "+" : ""}${round((primary.time_seconds?.mean ?? 0) - (baseline.time_seconds?.mean ?? 0), 1).toFixed(1)}`,
    tokens: `${(primary.tokens?.mean ?? 0) - (baseline.tokens?.mean ?? 0) >= 0 ? "+" : ""}${Math.round((primary.tokens?.mean ?? 0) - (baseline.tokens?.mean ?? 0))}`,
  };
  return runSummary;
}

export async function generateBenchmark(benchmarkDir, { skillName = "", skillPath = "" } = {}) {
  const results = await loadRunResults(benchmarkDir);
  const runSummary = aggregateResults(results);
  const runs = [];
  for (const [configuration, entries] of Object.entries(results)) {
    for (const result of entries) {
      runs.push({
        eval_id: result.eval_id,
        configuration,
        run_number: result.run_number,
        result: {
          pass_rate: result.pass_rate,
          passed: result.passed,
          failed: result.failed,
          total: result.total,
          time_seconds: result.time_seconds,
          tokens: result.tokens,
          tool_calls: result.tool_calls,
          errors: result.errors,
        },
        expectations: result.expectations,
        notes: result.notes,
      });
    }
  }
  const evalIds = [...new Set(runs.map((run) => run.eval_id))].sort((a, b) => a - b);
  return {
    metadata: {
      skill_name: skillName || "<skill-name>",
      skill_path: skillPath || "<path/to/skill>",
      executor_model: "<model-name>",
      analyzer_model: "<model-name>",
      timestamp: new Date().toISOString().replace(/\.\d{3}Z$/, "Z"),
      evals_run: evalIds,
      runs_per_configuration: 3,
    },
    runs,
    run_summary: runSummary,
    notes: [],
  };
}

export function generateMarkdown(benchmark) {
  const metadata = benchmark.metadata ?? {};
  const runSummary = benchmark.run_summary ?? {};
  const configs = Object.keys(runSummary).filter((key) => key !== "delta");
  const configA = configs[0] ?? "config_a";
  const configB = configs[1] ?? "config_b";
  const aSummary = runSummary[configA] ?? {};
  const bSummary = runSummary[configB] ?? {};
  const delta = runSummary.delta ?? {};
  const labelA = configA.replaceAll("_", " ").replace(/\b\w/g, (char) => char.toUpperCase());
  const labelB = configB.replaceAll("_", " ").replace(/\b\w/g, (char) => char.toUpperCase());
  const lines = [
    `# Skill Benchmark: ${metadata.skill_name ?? "<skill-name>"}`,
    "",
    `**Model**: ${metadata.executor_model ?? "<model-name>"}`,
    `**Date**: ${metadata.timestamp ?? ""}`,
    `**Evals**: ${(metadata.evals_run ?? []).join(", ")} (${metadata.runs_per_configuration ?? 0} runs each per configuration)`,
    "",
    "## Summary",
    "",
    `| Metric | ${labelA} | ${labelB} | Delta |`,
    "|--------|------------|---------------|-------|",
    `| Pass Rate | ${(aSummary.pass_rate?.mean ?? 0) * 100}% ± ${(aSummary.pass_rate?.stddev ?? 0) * 100}% | ${(bSummary.pass_rate?.mean ?? 0) * 100}% ± ${(bSummary.pass_rate?.stddev ?? 0) * 100}% | ${delta.pass_rate ?? "—"} |`,
    `| Time | ${(aSummary.time_seconds?.mean ?? 0).toFixed?.(1) ?? aSummary.time_seconds?.mean ?? 0}s ± ${(aSummary.time_seconds?.stddev ?? 0).toFixed?.(1) ?? aSummary.time_seconds?.stddev ?? 0}s | ${(bSummary.time_seconds?.mean ?? 0).toFixed?.(1) ?? bSummary.time_seconds?.mean ?? 0}s ± ${(bSummary.time_seconds?.stddev ?? 0).toFixed?.(1) ?? bSummary.time_seconds?.stddev ?? 0}s | ${delta.time_seconds ?? "—"}s |`,
    `| Tokens | ${Math.round(aSummary.tokens?.mean ?? 0)} ± ${Math.round(aSummary.tokens?.stddev ?? 0)} | ${Math.round(bSummary.tokens?.mean ?? 0)} ± ${Math.round(bSummary.tokens?.stddev ?? 0)} | ${delta.tokens ?? "—"} |`,
  ];
  if (Array.isArray(benchmark.notes) && benchmark.notes.length > 0) {
    lines.push("", "## Notes", "");
    for (const note of benchmark.notes) {
      lines.push(`- ${note}`);
    }
  }
  return lines.join("\n");
}

function parseArgs(args) {
  const flags = { _: [] };
  for (let i = 0; i < args.length; i += 1) {
    const item = args[i];
    if (!item.startsWith("-")) {
      flags._.push(item);
      continue;
    }
    const key = item.startsWith("--") ? item.slice(2) : item.slice(1);
    const next = args[i + 1];
    if (!next || next.startsWith("-")) {
      flags[key] = true;
      continue;
    }
    flags[key] = next;
    i += 1;
  }
  return flags;
}

async function main(argv = process.argv) {
  const args = parseArgs(argv.slice(2));
  const benchmarkDir = args._[0];
  if (!benchmarkDir) {
    throw new Error("Usage: node aggregate_benchmark.mjs <benchmark_dir> [--skill-name name] [--skill-path path] [-o output.json]");
  }
  const outputJsonPath = path.resolve(args.output ?? args.o ?? path.join(benchmarkDir, "benchmark.json"));
  const benchmark = await generateBenchmark(path.resolve(benchmarkDir), {
    skillName: args["skill-name"] ?? "",
    skillPath: args["skill-path"] ?? "",
  });
  await writeFile(outputJsonPath, `${JSON.stringify(benchmark, null, 2)}\n`);
  await writeFile(outputJsonPath.replace(/\.json$/i, ".md"), `${generateMarkdown(benchmark)}\n`);
  process.stdout.write(`Generated: ${outputJsonPath}\n`);
  process.stdout.write(`Generated: ${outputJsonPath.replace(/\.json$/i, ".md")}\n`);
}

const entryPath = fileURLToPath(import.meta.url);
if (process.argv[1] && path.resolve(process.argv[1]) === entryPath) {
  main().catch((error) => {
    process.stderr.write(`${error instanceof Error ? error.message : String(error)}\n`);
    process.exitCode = 1;
  });
}
