/**
 * Re-run Stage 1 annotations on orphan-word-patched paragraphs.
 *
 * Loads paragraphs that had orphan words restored, runs all 3 Stage 1 models
 * on the PATCHED text, and saves to a separate annotation file.
 * Original annotations in stage1.jsonl are NOT modified.
 *
 * Usage:
 *   bun ts/scripts/rerun-orphan-stage1.ts [--concurrency 60]
 *
 * Input:
 *   data/paragraphs/training.patched.jsonl  — patched paragraph text
 *   data/paragraphs/patches/orphan-word-patches.jsonl — patch records (for ID filtering)
 *
 * Output:
 *   data/annotations/stage1-orphan-rerun.jsonl — new annotations (separate file)
 */
import { readJsonl, readJsonlRaw, appendJsonl } from "../src/lib/jsonl.ts";
import { Paragraph } from "@sec-cybert/schemas/paragraph.ts";
import { STAGE1_MODELS } from "../src/lib/openrouter.ts";
import { annotateParagraph, type AnnotateOpts } from "../src/label/annotate.ts";
import { PROMPT_VERSION } from "../src/label/prompts.ts";
import { v4 as uuidv4 } from "uuid";
import { mkdir } from "node:fs/promises";
import { existsSync, readFileSync } from "node:fs";
import pLimit from "p-limit";

// ── Args ────────────────────────────────────────────────────────────────
const args = process.argv.slice(2);
function flag(name: string): string | undefined {
  const idx = args.indexOf(`--${name}`);
  return idx === -1 ? undefined : args[idx + 1];
}
const CONCURRENCY = parseInt(flag("concurrency") ?? "60", 10);

const DATA_DIR = new URL("../../data", import.meta.url).pathname;
const TRAINING_PATH = `${DATA_DIR}/paragraphs/training.patched.jsonl`;
const PATCHES_PATH = `${DATA_DIR}/paragraphs/patches/orphan-word-patches.jsonl`;
const OUTPUT_DIR = `${DATA_DIR}/annotations`;
const OUTPUT_PATH = `${OUTPUT_DIR}/stage1-orphan-rerun.jsonl`;

// ── Main ────────────────────────────────────────────────────────────────
async function main() {
  if (!existsSync(OUTPUT_DIR)) await mkdir(OUTPUT_DIR, { recursive: true });

  // Load orphan-word patch IDs
  console.error("Loading orphan-word patch IDs...");
  const patchIds = new Set<string>();
  for (const line of readFileSync(PATCHES_PATH, "utf-8").split("\n")) {
    if (!line.trim()) continue;
    const rec = JSON.parse(line) as { id: string };
    patchIds.add(rec.id);
  }
  console.error(`  ${patchIds.size} patched paragraph IDs`);

  // Load patched training data, filter to orphan-word paragraphs only
  console.error(`Loading patched paragraphs from ${TRAINING_PATH}...`);
  const { records: allParagraphs, skipped } = await readJsonl(TRAINING_PATH, Paragraph);
  if (skipped > 0) console.error(`  ⚠ Skipped ${skipped} invalid lines`);

  const paragraphs = allParagraphs.filter((p) => patchIds.has(p.id));
  console.error(`  ${paragraphs.length} orphan-word paragraphs in training set`);
  console.error(`  Models: ${STAGE1_MODELS.join(", ")}`);
  console.error(`  Prompt: ${PROMPT_VERSION}`);
  console.error(`  Concurrency: ${CONCURRENCY}`);

  const totalJobs = paragraphs.length * STAGE1_MODELS.length;
  console.error(`  Total annotations needed: ${totalJobs.toLocaleString()}`);

  // Load existing results for resume
  const doneKeys = new Set<string>();
  let resumedCost = 0;
  if (existsSync(OUTPUT_PATH)) {
    const { records: existing } = await readJsonlRaw(OUTPUT_PATH);
    for (const rec of existing) {
      const r = rec as { paragraphId?: string; provenance?: { modelId?: string; costUsd?: number } };
      if (r.paragraphId && r.provenance?.modelId) {
        doneKeys.add(`${r.paragraphId}|${r.provenance.modelId}`);
        resumedCost += r.provenance.costUsd ?? 0;
      }
    }
    if (doneKeys.size > 0) {
      console.error(`  Resuming: ${doneKeys.size} already done ($${resumedCost.toFixed(2)}), ${totalJobs - doneKeys.size} remaining`);
    }
  }

  if (doneKeys.size >= totalJobs) {
    console.error("  All annotations already complete!");
    return;
  }

  // Build job list
  type Job = { paragraph: Paragraph; modelId: string };
  const jobs: Job[] = [];
  for (const paragraph of paragraphs) {
    for (const modelId of STAGE1_MODELS) {
      if (!doneKeys.has(`${paragraph.id}|${modelId}`)) {
        jobs.push({ paragraph, modelId });
      }
    }
  }
  console.error(`  Jobs to run: ${jobs.length.toLocaleString()}\n`);

  // Run with concurrency limiter
  const runId = uuidv4();
  const limit = pLimit(CONCURRENCY);
  let completed = doneKeys.size;
  let failed = 0;
  let sessionCost = 0;
  const startTime = Date.now();

  // Progress logging
  const logInterval = setInterval(() => {
    const elapsed = (Date.now() - startTime) / 1000;
    const done = completed - doneKeys.size;
    const rate = done / elapsed;
    const remaining = totalJobs - completed;
    const eta = rate > 0 ? remaining / rate : Infinity;
    const etaMin = Math.floor(eta / 60);
    const etaSec = Math.round(eta % 60);
    process.stderr.write(
      `\x1b[2K\r  ${completed.toLocaleString()}/${totalJobs.toLocaleString()} (${((completed / totalJobs) * 100).toFixed(1)}%)` +
        `  $${(resumedCost + sessionCost).toFixed(4)}` +
        `  ${rate.toFixed(1)}/s` +
        `  ETA ${etaMin}m${etaSec.toString().padStart(2, "0")}s` +
        `  ${failed} failed`,
    );
  }, 2000);

  const tasks = jobs.map((job) =>
    limit(async () => {
      const opts: AnnotateOpts = {
        modelId: job.modelId,
        stage: "stage1",
        runId,
        promptVersion: PROMPT_VERSION,
        reasoningEffort: "low",
      };

      try {
        const ann = await annotateParagraph(job.paragraph, opts);
        await appendJsonl(OUTPUT_PATH, ann);
        sessionCost += ann.provenance.costUsd;
        completed++;
      } catch (error) {
        failed++;
        const msg = error instanceof Error ? error.message : String(error);
        console.error(`\n  ✖ ${job.modelId} × ${job.paragraph.id}: ${msg}`);
      }
    }),
  );

  await Promise.all(tasks);
  clearInterval(logInterval);

  const elapsed = ((Date.now() - startTime) / 1000).toFixed(0);
  console.error(
    `\n\n  ═══ ORPHAN WORD RE-ANNOTATION COMPLETE ═══` +
      `\n  Annotations: ${completed.toLocaleString()}/${totalJobs.toLocaleString()}` +
      `\n  Failed: ${failed}` +
      `\n  Session cost: $${sessionCost.toFixed(4)}` +
      `\n  Total cost: $${(resumedCost + sessionCost).toFixed(4)}` +
      `\n  Wall time: ${elapsed}s` +
      `\n  Output: ${OUTPUT_PATH}`,
  );

  if (failed > 0) {
    console.error(`\n  ⚠ ${failed} failures — re-run this script to retry them.`);
  }
}

main().catch((err) => {
  console.error(err);
  process.exit(1);
});
