#!/usr/bin/env node

import os from "node:os";
import path from "node:path";
import { mkdir, readFile, writeFile } from "node:fs/promises";
import { fileURLToPath } from "node:url";
import { generateHtml } from "./generate_report.mjs";
import { improveDescription } from "./improve_description.mjs";
import { findProjectRoot, openInBrowser } from "./claude_cli.mjs";
import { runEval } from "./run_eval.mjs";
import { parseSkillFile, resolveSkillPath } from "./skill_metadata.mjs";

function seededShuffle(values, seed) {
  let state = seed >>> 0;
  const random = () => {
    state = (1664525 * state + 1013904223) >>> 0;
    return state / 0x100000000;
  };
  const items = [...values];
  for (let i = items.length - 1; i > 0; i -= 1) {
    const j = Math.floor(random() * (i + 1));
    [items[i], items[j]] = [items[j], items[i]];
  }
  return items;
}

export function splitEvalSet(evalSet, holdout, seed = 42) {
  const trigger = seededShuffle(evalSet.filter((item) => item.should_trigger), seed);
  const noTrigger = seededShuffle(evalSet.filter((item) => !item.should_trigger), seed + 1);
  const triggerTestCount = Math.max(1, Math.floor(trigger.length * holdout));
  const noTriggerTestCount = Math.max(1, Math.floor(noTrigger.length * holdout));
  return {
    trainSet: [...trigger.slice(triggerTestCount), ...noTrigger.slice(noTriggerTestCount)],
    testSet: [...trigger.slice(0, triggerTestCount), ...noTrigger.slice(0, noTriggerTestCount)],
  };
}

function summarizeResults(results) {
  const passed = results.filter((item) => item.pass).length;
  return { passed, failed: results.length - passed, total: results.length };
}

function aggregateRuns(results) {
  let correct = 0;
  let total = 0;
  for (const result of results) {
    const runs = Number(result.runs ?? 0);
    const triggers = Number(result.triggers ?? 0);
    total += runs;
    correct += result.should_trigger ? triggers : runs - triggers;
  }
  return { correct, total };
}

function printEvalStats(label, results, elapsedSeconds = 0) {
  const pos = results.filter((item) => item.should_trigger);
  const neg = results.filter((item) => !item.should_trigger);
  const tp = pos.reduce((sum, item) => sum + Number(item.triggers ?? 0), 0);
  const posRuns = pos.reduce((sum, item) => sum + Number(item.runs ?? 0), 0);
  const fn = posRuns - tp;
  const fp = neg.reduce((sum, item) => sum + Number(item.triggers ?? 0), 0);
  const negRuns = neg.reduce((sum, item) => sum + Number(item.runs ?? 0), 0);
  const tn = negRuns - fp;
  const total = tp + tn + fp + fn;
  const precision = (tp + fp) > 0 ? tp / (tp + fp) : 1;
  const recall = (tp + fn) > 0 ? tp / (tp + fn) : 1;
  const accuracy = total > 0 ? (tp + tn) / total : 0;
  process.stderr.write(`${label}: ${tp + tn}/${total} correct, precision=${Math.round(precision * 100)}% recall=${Math.round(recall * 100)}% accuracy=${Math.round(accuracy * 100)}% (${elapsedSeconds.toFixed(1)}s)\n`);
  for (const result of results) {
    process.stderr.write(`  [${result.pass ? "PASS" : "FAIL"}] rate=${result.triggers}/${result.runs} expected=${result.should_trigger}: ${String(result.query).slice(0, 60)}\n`);
  }
}

export async function runLoop({
  evalSet,
  skillPath,
  descriptionOverride = null,
  numWorkers = 10,
  timeout = 30,
  maxIterations = 5,
  runsPerQuery = 3,
  triggerThreshold = 0.5,
  holdout = 0.4,
  model,
  verbose = false,
  liveReportPath = null,
  logDir = null,
}) {
  const { name, description: originalDescription, content } = await parseSkillFile(skillPath);
  let currentDescription = descriptionOverride ?? originalDescription;
  const { trainSet, testSet } = holdout > 0 ? splitEvalSet(evalSet, holdout) : { trainSet: evalSet, testSet: [] };
  const history = [];
  let exitReason = "unknown";

  for (let iteration = 1; iteration <= maxIterations; iteration += 1) {
    if (verbose) {
      process.stderr.write(`\n${"=".repeat(60)}\nIteration ${iteration}/${maxIterations}\nDescription: ${currentDescription}\n${"=".repeat(60)}\n`);
    }
    const allQueries = [...trainSet, ...testSet];
    const evalStart = Date.now();
    const allResults = await runEval(
      allQueries,
      name,
      currentDescription,
      numWorkers,
      timeout,
      findProjectRoot(),
      runsPerQuery,
      triggerThreshold,
      model ?? null,
    );
    const evalElapsed = (Date.now() - evalStart) / 1000;
    const trainQuerySet = new Set(trainSet.map((item) => item.query));
    const trainResultList = allResults.results.filter((item) => trainQuerySet.has(item.query));
    const testResultList = allResults.results.filter((item) => !trainQuerySet.has(item.query));
    const trainSummary = summarizeResults(trainResultList);
    const testSummary = testSet.length > 0 ? summarizeResults(testResultList) : null;
    history.push({
      iteration,
      description: currentDescription,
      train_passed: trainSummary.passed,
      train_failed: trainSummary.failed,
      train_total: trainSummary.total,
      train_results: trainResultList,
      test_passed: testSummary?.passed ?? null,
      test_failed: testSummary?.failed ?? null,
      test_total: testSummary?.total ?? null,
      test_results: testSummary ? testResultList : null,
      passed: trainSummary.passed,
      failed: trainSummary.failed,
      total: trainSummary.total,
      results: trainResultList,
    });
    if (liveReportPath) {
      const partialOutput = {
        original_description: originalDescription,
        best_description: currentDescription,
        best_score: "in progress",
        iterations_run: history.length,
        holdout,
        train_size: trainSet.length,
        test_size: testSet.length,
        history,
      };
      await writeFile(liveReportPath, generateHtml(partialOutput, { autoRefresh: true, skillName: name }));
    }
    if (verbose) {
      printEvalStats("Train", trainResultList, evalElapsed);
      if (testSummary) {
        printEvalStats("Test ", testResultList, 0);
      }
    }
    if (trainSummary.failed === 0) {
      exitReason = `all_passed (iteration ${iteration})`;
      break;
    }
    if (iteration === maxIterations) {
      exitReason = `max_iterations (${maxIterations})`;
      break;
    }
    const blindedHistory = history.map((item) => Object.fromEntries(Object.entries(item).filter(([key]) => !key.startsWith("test_"))));
    currentDescription = await improveDescription({
      skillName: name,
      skillContent: content,
      currentDescription,
      evalResults: { results: trainResultList, summary: trainSummary },
      history: blindedHistory,
      model,
      logDir,
      iteration,
    });
  }

  const best = testSet.length > 0
    ? history.reduce((bestItem, item) => ((item.test_passed ?? -1) > (bestItem?.test_passed ?? -1) ? item : bestItem), null)
    : history.reduce((bestItem, item) => ((item.train_passed ?? -1) > (bestItem?.train_passed ?? -1) ? item : bestItem), null);
  const bestScore = testSet.length > 0 ? `${best?.test_passed}/${best?.test_total}` : `${best?.train_passed}/${best?.train_total}`;
  return {
    exit_reason: exitReason,
    original_description: originalDescription,
    best_description: best?.description ?? currentDescription,
    best_score: bestScore,
    best_train_score: `${best?.train_passed ?? 0}/${best?.train_total ?? 0}`,
    best_test_score: testSet.length > 0 ? `${best?.test_passed ?? 0}/${best?.test_total ?? 0}` : null,
    final_description: currentDescription,
    iterations_run: history.length,
    holdout,
    train_size: trainSet.length,
    test_size: testSet.length,
    history,
  };
}

function parseArgs(args) {
  const flags = {};
  for (let i = 0; i < args.length; i += 1) {
    const value = args[i];
    if (!value.startsWith("--")) continue;
    const key = value.slice(2);
    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));
  if (args.help || args.h) {
    printHelp();
    return;
  }
  if (!args["skill-path"] || !args["eval-set"]) {
    printHelp();
    process.exitCode = 1;
    return;
  }
  const evalSet = JSON.parse(await readFile(args["eval-set"], "utf8"));
  const skillPath = resolveSkillPath(args["skill-path"]);
  const { name } = await parseSkillFile(skillPath);
  const reportMode = args.report ?? "auto";
  let liveReportPath = null;
  if (reportMode !== "none") {
    liveReportPath = reportMode === "auto"
      ? path.join(os.tmpdir(), `skill_description_report_${path.basename(skillPath)}_${Date.now()}.html`)
      : path.resolve(reportMode);
    await writeFile(liveReportPath, "<html><body><h1>Starting optimization loop...</h1><meta http-equiv='refresh' content='5'></body></html>");
    await openInBrowser(liveReportPath);
  }
  const resultsDir = args["results-dir"] ? path.join(path.resolve(args["results-dir"]), new Date().toISOString().replaceAll(":", "-")) : null;
  if (resultsDir) {
    await mkdir(resultsDir, { recursive: true });
  }
  const output = await runLoop({
    evalSet,
    skillPath,
    descriptionOverride: args.description ?? null,
    numWorkers: Number(args["num-workers"] ?? 10),
    timeout: Number(args.timeout ?? 30),
    maxIterations: Number(args["max-iterations"] ?? 5),
    runsPerQuery: Number(args["runs-per-query"] ?? 3),
    triggerThreshold: Number(args["trigger-threshold"] ?? 0.5),
    holdout: Number(args.holdout ?? 0.4),
    model: args.model,
    verbose: Boolean(args.verbose),
    liveReportPath,
    logDir: resultsDir ? path.join(resultsDir, "logs") : null,
  });
  const outputJson = `${JSON.stringify(output, null, 2)}\n`;
  process.stdout.write(outputJson);
  if (resultsDir) {
    await writeFile(path.join(resultsDir, "results.json"), outputJson);
  }
  if (liveReportPath) {
    await writeFile(liveReportPath, generateHtml(output, { autoRefresh: false, skillName: name }));
  }
  if (resultsDir) {
    await writeFile(path.join(resultsDir, "report.html"), `${generateHtml(output, { autoRefresh: false, skillName: name })}\n`);
  }
}

function printHelp() {
  process.stdout.write(
    [
      "Usage: node run_loop.mjs --skill-path <skill-dir|SKILL.md> --eval-set <eval-set.json> [options]",
      "",
      "Options:",
      "  --description <text>         Override the initial description",
      "  --num-workers <n>            Number of concurrent workers (default: 10)",
      "  --timeout <seconds>          Timeout per query (default: 30)",
      "  --max-iterations <n>         Maximum optimization iterations (default: 5)",
      "  --runs-per-query <n>         Number of runs per query (default: 3)",
      "  --trigger-threshold <ratio>  Pass threshold for positive prompts (default: 0.5)",
      "  --holdout <ratio>            Fraction reserved for test set (default: 0.4)",
      "  --report <auto|none|path>    Live report mode or output path (default: auto)",
      "  --results-dir <dir>          Directory for loop artifacts",
      "  --model <id>                 Optional model override",
      "  --verbose                    Print per-iteration stats",
      "  -h, --help                   Show this help",
      "",
      "The --skill-path argument may point to either the skill directory or its SKILL.md file.",
      "",
    ].join("\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;
  });
}
