import pLimit from "p-limit";
import { readJsonl } from "./lib/jsonl.ts";
import { Paragraph } from "@sec-cybert/schemas/paragraph.ts";
import { Annotation } from "@sec-cybert/schemas/annotation.ts";
import { STAGE1_MODEL, STAGE1_RUNS, BENCHMARK_MODELS } from "./lib/openrouter.ts";
import { runBatch } from "./label/batch.ts";
import { runGoldenBatch } from "./label/golden.ts";
import { computeConsensus } from "./label/consensus.ts";
import { judgeParagraph, annotateParagraph, reEvalParagraph } from "./label/annotate.ts";
import { appendJsonl, readJsonlRaw } from "./lib/jsonl.ts";
import { v4 as uuidv4 } from "uuid";
import { PROMPT_VERSION } from "./label/prompts.ts";
import { extract10K, extract8K, reparse10K, reparse8K, mergeTrainingData } from "./extract/pipeline.ts";

const DATA = "../data";

function usage(): never {
  console.error(`Usage: bun sec <command> [options]

Commands:
  extract:10k      [--start-date YYYY-MM-DD] [--end-date YYYY-MM-DD] [--limit N]
  extract:8k       [--start-date YYYY-MM-DD] [--end-date YYYY-MM-DD] [--limit N]
  extract:reparse     Re-parse cached 10-K HTML files with current parser (no network)
  extract:reparse-8k  Re-parse cached 8-K HTML files with current parser (no network)
  extract:merge       Merge 10-K + 8-K, remove truncated filings, dedup → training.jsonl
  label:annotate   [--model <id>] [--run N] [--output-dir DIR] [--paragraphs PATH] [--limit N] [--concurrency N]
  label:annotate-all [--output-dir DIR] [--paragraphs PATH] [--limit N] [--concurrency N]  (Grok ×3 self-consistency)
  label:consensus  [--input-dir DIR]
  label:judge [--concurrency N]
  label:golden [--paragraphs <path>] [--limit N] [--delay N] [--concurrency N]  (Opus via Agent SDK)
  label:bench-holdout --model <id> [--concurrency N] [--limit N]   (benchmark model on holdout)
  label:bench-holdout-all [--concurrency N] [--limit N]            (all BENCHMARK_MODELS on holdout)
  label:bench-holdout-v35 --model <id> [--concurrency N] [--limit N]  (v3.5 re-run on confusion-axis holdout)
  label:golden-v35 [--limit N] [--delay N] [--concurrency N]        (Opus v3.5 re-run on confusion-axis holdout)
  label:reeval --model <id> [--concurrency N] [--limit N]          (re-evaluate flagged Stage 1 paragraphs)
  label:cost`);
  process.exit(1);
}

const [command, ...rest] = process.argv.slice(2);
if (!command) usage();

function flag(name: string): string | undefined {
  const idx = rest.indexOf(`--${name}`);
  if (idx === -1) return undefined;
  return rest[idx + 1];
}

function flagInt(name: string, fallback: number): number {
  const v = flag(name);
  return v !== undefined ? parseInt(v, 10) : fallback;
}

const PARAGRAPHS_PATH = `${DATA}/paragraphs/paragraphs.jsonl`;
const PARAGRAPHS_8K_PATH = `${DATA}/paragraphs/paragraphs-8k.jsonl`;
const SESSIONS_PATH = `${DATA}/metadata/sessions.jsonl`;

async function loadParagraphs(): Promise<Paragraph[]> {
  const { records, skipped } = await readJsonl(PARAGRAPHS_PATH, Paragraph);
  if (skipped > 0) process.stderr.write(`  ⚠ Skipped ${skipped} invalid paragraph lines\n`);
  if (records.length === 0) {
    process.stderr.write(`  ✖ No paragraphs found at ${PARAGRAPHS_PATH}\n`);
    process.exit(1);
  }
  process.stderr.write(`  Loaded ${records.length} paragraphs\n`);
  return records;
}

async function cmdAnnotate(): Promise<void> {
  const modelId = flag("model") ?? STAGE1_MODEL;
  const paragraphsPath = flag("paragraphs") ?? `${DATA}/paragraphs/paragraphs-clean.patched.jsonl`;
  const { records: paragraphs, skipped } = await readJsonl(paragraphsPath, Paragraph);
  if (skipped > 0) process.stderr.write(`  ⚠ Skipped ${skipped} invalid paragraph lines\n`);
  if (paragraphs.length === 0) {
    process.stderr.write(`  ✖ No paragraphs found at ${paragraphsPath}\n`);
    process.exit(1);
  }
  process.stderr.write(`  Loaded ${paragraphs.length} paragraphs from ${paragraphsPath}\n`);

  const modelShort = modelId.split("/")[1]!;
  const outputDir = flag("output-dir") ?? "v2-stage1";
  const runSuffix = flag("run") ? `.run${flag("run")}` : "";

  await runBatch(paragraphs, {
    modelId,
    stage: "stage1",
    outputPath: `${DATA}/annotations/${outputDir}/${modelShort}${runSuffix}.jsonl`,
    errorsPath: `${DATA}/annotations/${outputDir}/${modelShort}${runSuffix}-errors.jsonl`,
    sessionsPath: SESSIONS_PATH,
    concurrency: flagInt("concurrency", 60),
    limit: flag("limit") !== undefined ? flagInt("limit", 50) : undefined,
  });
}

async function cmdAnnotateAll(): Promise<void> {
  const paragraphsPath = flag("paragraphs") ?? `${DATA}/paragraphs/paragraphs-clean.patched.jsonl`;
  const { records: allParagraphs, skipped } = await readJsonl(paragraphsPath, Paragraph);
  if (skipped > 0) process.stderr.write(`  ⚠ Skipped ${skipped} invalid paragraph lines\n`);
  if (allParagraphs.length === 0) {
    process.stderr.write(`  ✖ No paragraphs found at ${paragraphsPath}\n`);
    process.exit(1);
  }
  process.stderr.write(`  Loaded ${allParagraphs.length} paragraphs from ${paragraphsPath}\n`);

  const concurrency = flagInt("concurrency", 60);
  const limit = flag("limit") !== undefined ? flagInt("limit", 50) : undefined;
  const outputDir = flag("output-dir") ?? "v2-stage1";
  const modelShort = STAGE1_MODEL.split("/")[1]!;

  process.stderr.write(`  Stage 1: ${STAGE1_MODEL} ×${STAGE1_RUNS} self-consistency → ${outputDir}/\n`);

  for (let run = 1; run <= STAGE1_RUNS; run++) {
    process.stderr.write(`\n  ═══ ${STAGE1_MODEL} run ${run}/${STAGE1_RUNS} ═══\n`);
    await runBatch(allParagraphs, {
      modelId: STAGE1_MODEL,
      stage: "stage1",
      outputPath: `${DATA}/annotations/${outputDir}/${modelShort}.run${run}.jsonl`,
      errorsPath: `${DATA}/annotations/${outputDir}/${modelShort}.run${run}-errors.jsonl`,
      sessionsPath: SESSIONS_PATH,
      concurrency,
      limit,
    });
  }
}

async function cmdConsensus(): Promise<void> {
  // Load all Stage 1 annotations (3 self-consistency runs)
  const allAnnotations: Map<string, Annotation[]> = new Map();
  const inputDir = flag("input-dir") ?? "v2-stage1";
  const modelShort = STAGE1_MODEL.split("/")[1]!;

  for (let run = 1; run <= STAGE1_RUNS; run++) {
    const path = `${DATA}/annotations/${inputDir}/${modelShort}.run${run}.jsonl`;
    const { records } = await readJsonl(path, Annotation);
    process.stderr.write(`  Loaded ${records.length} annotations from run ${run}\n`);
    for (const ann of records) {
      const existing = allAnnotations.get(ann.paragraphId) ?? [];
      existing.push(ann);
      allAnnotations.set(ann.paragraphId, existing);
    }
  }

  // Only process paragraphs with all 3 annotations
  let consensus = 0;
  let needsJudge = 0;
  const outputPath = `${DATA}/annotations/${inputDir}/consensus.jsonl`;

  for (const [paragraphId, anns] of allAnnotations) {
    if (anns.length !== 3) continue;

    const { result, needsJudge: needs } = computeConsensus(paragraphId, anns);
    await appendJsonl(outputPath, result);

    if (needs) needsJudge++;
    else consensus++;
  }

  const total = consensus + needsJudge;
  process.stderr.write(
    `\n  ✓ Consensus: ${consensus}/${total} (${((consensus / total) * 100).toFixed(1)}%) agreed\n` +
      `    ${needsJudge} paragraphs need Stage 2 judge\n`,
  );
}

async function cmdJudge(): Promise<void> {
  // Load paragraphs and consensus results needing judge
  const paragraphsPath = `${DATA}/paragraphs/paragraphs-clean.patched.jsonl`;
  const { records: paragraphs, skipped: pSkipped } = await readJsonl(paragraphsPath, Paragraph);
  if (pSkipped > 0) process.stderr.write(`  ⚠ Skipped ${pSkipped} invalid paragraph lines\n`);
  process.stderr.write(`  Loaded ${paragraphs.length} paragraphs\n`);
  const paragraphMap = new Map(paragraphs.map((p) => [p.id, p]));

  const judgeDir = flag("input-dir") ?? "v2-stage1";
  const consensusPath = `${DATA}/annotations/${judgeDir}/consensus.jsonl`;
  const { records: rawConsensus } = await readJsonlRaw(consensusPath);

  // Load all stage 1 annotations for lookup (3 self-consistency runs)
  const stage1Map: Map<string, Annotation[]> = new Map();
  const judgeModelShort = STAGE1_MODEL.split("/")[1]!;

  for (let run = 1; run <= STAGE1_RUNS; run++) {
    const { records } = await readJsonl(
      `${DATA}/annotations/${judgeDir}/${judgeModelShort}.run${run}.jsonl`,
      Annotation,
    );
    for (const ann of records) {
      const existing = stage1Map.get(ann.paragraphId) ?? [];
      existing.push(ann);
      stage1Map.set(ann.paragraphId, existing);
    }
  }

  const unresolvedIds: string[] = [];
  for (const raw of rawConsensus) {
    const r = raw as { paragraphId?: string; method?: string };
    if (
      r.method === "unresolved" &&
      r.paragraphId
    ) {
      unresolvedIds.push(r.paragraphId);
    }
  }

  // Check what's already judged
  const judgePath = `${DATA}/annotations/${judgeDir}/judge.jsonl`;
  const { records: existing } = await readJsonlRaw(judgePath);
  const judgedIds = new Set(
    existing
      .filter((r): r is { paragraphId: string } =>
        !!r && typeof r === "object" && "paragraphId" in r)
      .map((r) => r.paragraphId),
  );

  const toJudge = unresolvedIds.filter((id) => !judgedIds.has(id));
  const judgeModelId = flag("model") ?? "openai/gpt-5.4";
  const concurrency = flagInt("concurrency", 30);
  process.stderr.write(`  ${toJudge.length} paragraphs to judge (${judgedIds.size} already done) │ ${judgeModelId} │ concurrency=${concurrency}\n`);

  const runId = uuidv4();
  let processed = 0;
  let errored = 0;
  const limit_ = pLimit(concurrency);

  const tasks = toJudge.map((paragraphId) =>
    limit_(async () => {
      const paragraph = paragraphMap.get(paragraphId);
      if (!paragraph) return;

      const stage1Anns = stage1Map.get(paragraphId);
      if (!stage1Anns || stage1Anns.length < 3) return;

      const priorLabels = stage1Anns.map((a) => ({
        content_category: a.label.content_category,
        specificity_level: a.label.specificity_level,
        reasoning: a.label.reasoning,
      }));

      try {
        const judgeAnn = await judgeParagraph(paragraph, priorLabels, {
          runId,
          promptVersion: PROMPT_VERSION,
          modelId: judgeModelId,
        });
        await appendJsonl(judgePath, judgeAnn);
        processed++;

        if (processed % 10 === 0) {
          process.stderr.write(`  Judged ${processed}/${toJudge.length}\n`);
        }
      } catch (error) {
        errored++;
        process.stderr.write(
          `  ✖ Judge error for ${paragraphId}: ${error instanceof Error ? error.message : String(error)}\n`,
        );
      }
    }),
  );

  await Promise.all(tasks);
  process.stderr.write(`\n  ✓ Judged ${processed} paragraphs (${errored} errors)\n`);
}

async function cmdGolden(): Promise<void> {
  // Load holdout paragraph IDs (v2 default, override with --ids)
  const sampledIdsPath = flag("ids") ?? `${DATA}/gold/v2-holdout-ids.json`;
  const sampledIds = new Set<string>(
    JSON.parse(await import("node:fs/promises").then((fs) => fs.readFile(sampledIdsPath, "utf-8"))),
  );
  process.stderr.write(`  Loaded ${sampledIds.size} sampled IDs from ${sampledIdsPath}\n`);

  // Load patched paragraphs and filter to the human-labeled set
  const paragraphsPath = flag("paragraphs") ?? `${DATA}/paragraphs/paragraphs-clean.patched.jsonl`;
  const { records: allParagraphs, skipped } = await readJsonl(paragraphsPath, Paragraph);
  if (skipped > 0) process.stderr.write(`  ⚠ Skipped ${skipped} invalid paragraph lines\n`);

  const paragraphs = allParagraphs.filter((p) => sampledIds.has(p.id));
  process.stderr.write(`  Matched ${paragraphs.length}/${sampledIds.size} paragraphs from ${paragraphsPath}\n`);

  if (paragraphs.length === 0) {
    process.stderr.write("  ✖ No matching paragraphs found\n");
    process.exit(1);
  }

  const promptOnly = rest.includes("--prompt-only");
  const outputDir = promptOnly ? "v2-bench" : "golden";
  const outputFile = promptOnly ? "opus-4.6" : "opus";

  await runGoldenBatch(paragraphs, {
    outputPath: `${DATA}/annotations/${outputDir}/${outputFile}.jsonl`,
    errorsPath: `${DATA}/annotations/${outputDir}/${outputFile}-errors.jsonl`,
    limit: flag("limit") !== undefined ? flagInt("limit", 50) : undefined,
    delayMs: flag("delay") !== undefined ? flagInt("delay", 1000) : 1000,
    concurrency: flagInt("concurrency", 1),
    promptOnly,
  });
}

async function loadHoldoutParagraphs(idsPath?: string): Promise<Paragraph[]> {
  const sampledIdsPath = idsPath ?? `${DATA}/gold/v2-holdout-ids.json`;
  const sampledIds = new Set<string>(
    JSON.parse(await import("node:fs/promises").then((fs) => fs.readFile(sampledIdsPath, "utf-8"))),
  );
  process.stderr.write(`  Loaded ${sampledIds.size} holdout IDs from ${sampledIdsPath}\n`);

  const paragraphsPath = `${DATA}/paragraphs/paragraphs-clean.patched.jsonl`;
  const { records: allParagraphs, skipped } = await readJsonl(paragraphsPath, Paragraph);
  if (skipped > 0) process.stderr.write(`  ⚠ Skipped ${skipped} invalid paragraph lines\n`);

  const paragraphs = allParagraphs.filter((p) => sampledIds.has(p.id));
  process.stderr.write(`  Matched ${paragraphs.length}/${sampledIds.size} holdout paragraphs\n`);

  if (paragraphs.length === 0) {
    process.stderr.write("  ✖ No matching paragraphs found\n");
    process.exit(1);
  }
  return paragraphs;
}

async function cmdBenchHoldout(): Promise<void> {
  const modelId = flag("model");
  if (!modelId) {
    console.error("--model is required");
    process.exit(1);
  }
  const idsPath = flag("ids");
  const paragraphs = await loadHoldoutParagraphs(idsPath ?? undefined);
  const modelShort = modelId.split("/")[1]!;
  const outputDir = flag("output-dir") ?? "v2-bench";

  await runBatch(paragraphs, {
    modelId,
    stage: "benchmark",
    outputPath: `${DATA}/annotations/${outputDir}/${modelShort}.jsonl`,
    errorsPath: `${DATA}/annotations/${outputDir}/${modelShort}-errors.jsonl`,
    sessionsPath: SESSIONS_PATH,
    concurrency: flagInt("concurrency", 60),
    limit: flag("limit") !== undefined ? flagInt("limit", 50) : undefined,
  });
}

async function cmdBenchHoldoutAll(): Promise<void> {
  const idsPath = flag("ids");
  const paragraphs = await loadHoldoutParagraphs(idsPath ?? undefined);
  const concurrency = flagInt("concurrency", 60);
  const limit = flag("limit") !== undefined ? flagInt("limit", 50) : undefined;
  const outputDir = flag("output-dir") ?? "v2-bench";

  process.stderr.write(`  Running ${BENCHMARK_MODELS.length} benchmark models → ${outputDir}/\n`);

  for (const modelId of BENCHMARK_MODELS) {
    const modelShort = modelId.split("/")[1]!;
    process.stderr.write(`\n  ═══ ${modelId} ═══\n`);
    await runBatch(paragraphs, {
      modelId,
      stage: "benchmark",
      outputPath: `${DATA}/annotations/${outputDir}/${modelShort}.jsonl`,
      errorsPath: `${DATA}/annotations/${outputDir}/${modelShort}-errors.jsonl`,
      sessionsPath: SESSIONS_PATH,
      concurrency,
      limit,
    });
  }
}

async function loadConfusionAxisParagraphs(rerunFile?: string): Promise<Paragraph[]> {
  const rerunPath = rerunFile ?? `${DATA}/gold/holdout-rerun-v35.jsonl`;
  const { records: rerunRecords } = await readJsonlRaw(rerunPath);
  const rerunIds = new Set(
    rerunRecords
      .filter((r): r is { paragraphId: string } =>
        !!r && typeof r === "object" && "paragraphId" in r)
      .map((r) => r.paragraphId),
  );
  process.stderr.write(`  Loaded ${rerunIds.size} confusion-axis paragraph IDs\n`);

  const allHoldout = await loadHoldoutParagraphs();
  const paragraphs = allHoldout.filter((p) => rerunIds.has(p.id));
  process.stderr.write(`  Matched ${paragraphs.length} paragraphs for v3.5 re-run\n`);
  return paragraphs;
}

async function cmdGoldenV35(): Promise<void> {
  const paragraphs = await loadConfusionAxisParagraphs();

  if (paragraphs.length === 0) {
    process.stderr.write("  ✖ No confusion-axis paragraphs found\n");
    process.exit(1);
  }

  await runGoldenBatch(paragraphs, {
    outputPath: `${DATA}/annotations/golden-v35/opus.jsonl`,
    errorsPath: `${DATA}/annotations/golden-v35/opus-errors.jsonl`,
    limit: flag("limit") !== undefined ? flagInt("limit", 50) : undefined,
    delayMs: flag("delay") !== undefined ? flagInt("delay", 500) : 500,
    concurrency: flagInt("concurrency", 20),
  });
}

async function cmdBenchHoldoutV35(): Promise<void> {
  const modelId = flag("model");
  if (!modelId) {
    console.error("--model is required");
    process.exit(1);
  }

  const rerunFile = flag("rerun-file");
  const outputDir = flag("output-dir") ?? "bench-holdout-v35";
  const paragraphs = await loadConfusionAxisParagraphs(rerunFile ?? undefined);
  const modelShort = modelId.split("/")[1]!;
  await runBatch(paragraphs, {
    modelId,
    stage: "benchmark",
    outputPath: `${DATA}/annotations/${outputDir}/${modelShort}.jsonl`,
    errorsPath: `${DATA}/annotations/${outputDir}/${modelShort}-errors.jsonl`,
    sessionsPath: SESSIONS_PATH,
    concurrency: flagInt("concurrency", 60),
    limit: flag("limit") !== undefined ? flagInt("limit", 50) : undefined,
  });
}

async function cmdReeval(): Promise<void> {
  const modelId = flag("model");
  if (!modelId) {
    console.error("--model is required");
    process.exit(1);
  }

  // Load flagged paragraphs
  const correctionsPath = `${DATA}/annotations/stage1-corrections.jsonl`;
  const { records: corrections } = await readJsonlRaw(correctionsPath);
  const correctionMap = new Map<string, { reason: string }>();
  for (const r of corrections) {
    const rec = r as { paragraphId?: string; reason?: string };
    if (rec.paragraphId && rec.reason) {
      correctionMap.set(rec.paragraphId, { reason: rec.reason });
    }
  }
  process.stderr.write(`  Loaded ${correctionMap.size} flagged paragraphs from ${correctionsPath}\n`);

  // Load all paragraphs and filter to flagged ones
  const paragraphs = await loadParagraphs();
  const flaggedParagraphs = paragraphs.filter((p) => correctionMap.has(p.id));
  process.stderr.write(`  Matched ${flaggedParagraphs.length} paragraphs for re-evaluation\n`);

  const modelShort = modelId.split("/")[1]!;
  const outputPath = `${DATA}/annotations/stage2/reeval-${modelShort}.jsonl`;
  const errorsPath = `${DATA}/annotations/stage2/reeval-${modelShort}-errors.jsonl`;

  // Resume support
  const { records: existing } = await readJsonlRaw(outputPath);
  const doneIds = new Set(
    existing
      .filter((r): r is { paragraphId: string } =>
        !!r && typeof r === "object" && "paragraphId" in r)
      .map((r) => r.paragraphId),
  );

  let remaining = flaggedParagraphs.filter((p) => !doneIds.has(p.id));
  const limit = flag("limit") !== undefined ? flagInt("limit", 50) : undefined;
  if (limit !== undefined) remaining = remaining.slice(0, limit);

  process.stderr.write(`  ${remaining.length} paragraphs to re-evaluate (${doneIds.size} already done)\n`);

  const runId = uuidv4();
  const concurrency = flagInt("concurrency", 12);
  let processed = 0;
  let errored = 0;
  let totalCost = 0;

  // Process in batches respecting concurrency
  for (let i = 0; i < remaining.length; i += concurrency) {
    const batch = remaining.slice(i, i + concurrency);
    const results = await Promise.allSettled(
      batch.map(async (paragraph) => {
        const correction = correctionMap.get(paragraph.id)!;
        const reason = correction.reason as "materiality_language" | "spac" | "other";
        const ann = await reEvalParagraph(paragraph, {
          modelId,
          runId,
          reason,
          promptVersion: PROMPT_VERSION,
        });
        await appendJsonl(outputPath, ann);
        totalCost += ann.provenance.costUsd;
        processed++;
      }),
    );

    for (const r of results) {
      if (r.status === "rejected") {
        errored++;
        process.stderr.write(`  ✖ Error: ${r.reason}\n`);
      }
    }

    if (processed % 50 === 0 || i + concurrency >= remaining.length) {
      process.stderr.write(`  ${processed}/${remaining.length} re-evaluated (${errored} errors, $${totalCost.toFixed(4)} cost)\n`);
    }
  }

  process.stderr.write(`\n  ✓ Re-evaluation done: ${processed} processed, ${errored} errors, $${totalCost.toFixed(4)} total cost\n`);
}

async function cmdCost(): Promise<void> {
  const modelCosts: Record<string, { cost: number; count: number }> = {};
  const stageCosts: Record<string, { cost: number; count: number }> = {};

  // Stage 1 (Grok ×3 self-consistency runs)
  const costModelShort = STAGE1_MODEL.split("/")[1]!;
  for (let run = 1; run <= STAGE1_RUNS; run++) {
    const path = `${DATA}/annotations/v2-stage1/${costModelShort}.run${run}.jsonl`;
    try {
      const { records } = await readJsonl(path, Annotation);
      const cost = records.reduce((sum, a) => sum + a.provenance.costUsd, 0);
      const key = `${STAGE1_MODEL} (run ${run})`;
      modelCosts[key] = { cost, count: records.length };
      const stage = stageCosts["stage1"] ?? { cost: 0, count: 0 };
      stage.cost += cost;
      stage.count += records.length;
      stageCosts["stage1"] = stage;
    } catch {
      // Run file may not exist yet
    }
  }

  // Stage 2
  const judgePath = `${DATA}/annotations/stage2/judge.jsonl`;
  const { records: judgeRecords } = await readJsonl(judgePath, Annotation);
  const judgeCost = judgeRecords.reduce((sum, a) => sum + a.provenance.costUsd, 0);
  modelCosts["anthropic/claude-sonnet-4.6 (judge)"] = {
    cost: judgeCost,
    count: judgeRecords.length,
  };
  stageCosts["stage2"] = { cost: judgeCost, count: judgeRecords.length };

  // Print
  console.log("\n  Cost Summary");
  console.log("  " + "─".repeat(60));

  console.log("\n  By Model:");
  for (const [model, { cost, count }] of Object.entries(modelCosts)) {
    console.log(`    ${model.padEnd(45)} $${cost.toFixed(4)} (${count} annotations)`);
  }

  console.log("\n  By Stage:");
  let total = 0;
  for (const [stage, { cost, count }] of Object.entries(stageCosts)) {
    console.log(`    ${stage.padEnd(45)} $${cost.toFixed(4)} (${count} annotations)`);
    total += cost;
  }

  console.log(`\n  Total: $${total.toFixed(4)}`);
}

async function cmdExtract10K(): Promise<void> {
  await extract10K({
    outputPath: PARAGRAPHS_PATH,
    errorsPath: `${DATA}/extracted/item1c/errors.jsonl`,
    filingType: "10-K",
    startDate: flag("start-date") ?? "2023-12-15",
    endDate: flag("end-date") ?? "2025-12-31",
    limit: flag("limit") !== undefined ? flagInt("limit", 100) : undefined,
  });
}

async function cmdExtract8K(): Promise<void> {
  await extract8K({
    outputPath: PARAGRAPHS_8K_PATH,
    errorsPath: `${DATA}/extracted/item105/errors.jsonl`,
    filingType: "8-K",
    startDate: flag("start-date") ?? "2023-12-18",
    endDate: flag("end-date") ?? "2025-12-31",
    limit: flag("limit") !== undefined ? flagInt("limit", 100) : undefined,
  });
}

async function cmdReparse8K(): Promise<void> {
  await reparse8K({ outputPath: PARAGRAPHS_8K_PATH });
}

// Dispatch
switch (command) {
  case "extract:10k":
    await cmdExtract10K();
    break;
  case "extract:8k":
    await cmdExtract8K();
    break;
  case "extract:reparse":
    await reparse10K({ outputPath: PARAGRAPHS_PATH });
    break;
  case "extract:reparse-8k":
    await cmdReparse8K();
    break;
  case "extract:merge":
    await mergeTrainingData({
      tenKPath: PARAGRAPHS_PATH,
      eightKPath: PARAGRAPHS_8K_PATH,
      outputPath: `${DATA}/paragraphs/training.jsonl`,
    });
    break;
  case "label:annotate":
    await cmdAnnotate();
    break;
  case "label:annotate-all":
    await cmdAnnotateAll();
    break;
  case "label:consensus":
    await cmdConsensus();
    break;
  case "label:judge":
    await cmdJudge();
    break;
  case "label:golden":
    await cmdGolden();
    break;
  case "label:bench-holdout":
    await cmdBenchHoldout();
    break;
  case "label:bench-holdout-all":
    await cmdBenchHoldoutAll();
    break;
  case "label:bench-holdout-v35":
    await cmdBenchHoldoutV35();
    break;
  case "label:golden-v35":
    await cmdGoldenV35();
    break;
  case "label:reeval":
    await cmdReeval();
    break;
  case "label:cost":
    await cmdCost();
    break;
  default:
    console.error(`Unknown command: ${command}`);
    usage();
}
