//  File:    rate.ts
//  Path:    apps/worker/src/scripts/rate.ts
//  Project: AI Risk Index — airiskindex.io
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//
//  Description: One-shot parallel rating run with per-model resume via state file.

import { prisma } from "@airiskindex/db";
import { RATER_MODELS, RATER_PROMPT_VERSION } from "../config";
import {
  ingestBatchResults,
  isBatchComplete,
  submitRatingBatch,
  type RatingTask,
} from "../raters/batch";
import { ratingJobId } from "../raters/job-id";

// One-shot rating run: submit a Message Batch per model for every task that
// lacks ratings under the current prompt version, poll to completion, ingest.
// Models run in PARALLEL (batches are independent server-side queues).
// Usage: tsx src/scripts/rate.ts [--limit N] [--poll-seconds S]
// Idempotent & resumable: already-rated tasks are excluded, and an in-flight
// batch per model is remembered in a state file so a relaunch resumes polling
// instead of resubmitting (a resubmit would double the spend).

import { readFileSync, writeFileSync } from "node:fs";

function argValue(flag: string): string | undefined {
  const index = process.argv.indexOf(flag);
  return index >= 0 ? process.argv[index + 1] : undefined;
}

const limit = Number(argValue("--limit") ?? "0") || undefined;
const pollSeconds = Number(argValue("--poll-seconds") ?? "60") || 60;
const STATE_PATH = process.env.RATE_STATE ?? "/tmp/aix-batch-state.json";

function readState(): Record<string, string> {
  try {
    return JSON.parse(readFileSync(STATE_PATH, "utf8")) as Record<string, string>;
  } catch {
    return {};
  }
}

function writeState(state: Record<string, string>): void {
  writeFileSync(STATE_PATH, JSON.stringify(state, null, 2));
}

const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));

async function runForModel(model: string): Promise<void> {
  const tasks = await prisma.task.findMany({
    where: {
      ratings: { none: { model, promptVersion: RATER_PROMPT_VERSION } },
      // exclude demo-seeded synthetic tasks (real O*NET task IDs are numeric)
      NOT: { id: { contains: "-T" } },
    },
    include: { occupation: { select: { title: true } } },
    orderBy: { id: "asc" },
    ...(limit ? { take: limit } : {}),
  });
  if (tasks.length === 0) {
    console.log(`[${model}] nothing to rate`);
    return;
  }

  const ratingTasks: RatingTask[] = tasks.map((task) => ({
    taskId: task.id,
    occupationTitle: task.occupation.title,
    statement: task.statement,
  }));
  const jobIdToTaskId = new Map(
    ratingTasks.map((task) => [ratingJobId(task.taskId, model, RATER_PROMPT_VERSION), task.taskId]),
  );

  let batchId = readState()[model];
  if (batchId) {
    console.log(`[${model}] resuming in-flight batch ${batchId}; polling every ${pollSeconds}s`);
  } else {
    console.log(`[${model}] submitting batch of ${ratingTasks.length} tasks…`);
    batchId = await submitRatingBatch(ratingTasks, model);
    writeState({ ...readState(), [model]: batchId });
    console.log(`[${model}] batch ${batchId} submitted; polling every ${pollSeconds}s`);
  }

  for (;;) {
    await sleep(pollSeconds * 1000);
    if (await isBatchComplete(batchId)) break;
    console.log(`[${model}] batch ${batchId} still processing…`);
  }

  const { ingested, failed } = await ingestBatchResults(batchId, model, jobIdToTaskId);
  const state = readState();
  delete state[model];
  writeState(state);
  console.log(`[${model}] batch ${batchId} done: ${ingested} dimension ratings ingested, ${failed} requests failed`);
}

async function main(): Promise<void> {
  if (RATER_MODELS.length < 2) {
    throw new Error("RATER_MODELS must list ≥2 models (multi-model rating, CLAUDE.md §6)");
  }
  console.log(
    `rating run — prompt ${RATER_PROMPT_VERSION}, models: ${RATER_MODELS.join(", ")}${limit ? `, limit ${limit}` : " (all unrated tasks)"}`,
  );
  await Promise.all(RATER_MODELS.map((model) => runForModel(model)));
  console.log("rating run complete — next: pnpm score:recompute");
}

main()
  .catch((error) => {
    console.error(error);
    process.exitCode = 1;
  })
  .finally(() => prisma.$disconnect());