//  File:    batch.ts
//  Path:    apps/worker/src/raters/batch.ts
//  Project: AI Risk Index — airiskindex.io
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//
//  Description: Anthropic Message Batches integration: submit, poll, ingest with audit trail.

import { readFileSync } from "node:fs";
import Anthropic from "@anthropic-ai/sdk";
import { prisma } from "@airiskindex/db";
import { RATER_PROMPT_VERSION } from "../config";
import { ratingJobId } from "./job-id";
import { RATING_OUTPUT_JSON_SCHEMA, ratingResponseSchema } from "./schema";

// Batch rating flow (docs/research/03-llm-rater-api.md): one Message Batches
// request per task × model, custom_id = deterministic job ID, 1h-cached rubric
// in the system block, structured JSON output. 50% batch discount; results
// retrievable for 29 days.

const anthropic = new Anthropic();

const rubric = readFileSync(new URL("./prompts/task-rating-v1.md", import.meta.url), "utf8");

export interface RatingTask {
  taskId: string;
  occupationTitle: string;
  statement: string;
}

export async function submitRatingBatch(tasks: RatingTask[], model: string): Promise<string> {
  const batch = await anthropic.messages.batches.create({
    requests: tasks.map((task) => ({
      custom_id: ratingJobId(task.taskId, model, RATER_PROMPT_VERSION),
      params: {
        model,
        max_tokens: 3000,
        system: [
          {
            type: "text" as const,
            text: rubric,
            cache_control: { type: "ephemeral" as const },
          },
        ],
        messages: [
          {
            role: "user" as const,
            content: `Occupation: ${task.occupationTitle}\nTask statement: ${task.statement}\n\nRate this task per the rubric.`,
          },
        ],
        // Structured outputs (GA) + thinking disabled (Sonnet 5 thinks by
        // default and thinking bills as output). Untyped extension so the code
        // compiles across SDK versions; see docs/research/03-llm-rater-api.md.
        ...({
          output_config: {
            format: { type: "json_schema", schema: RATING_OUTPUT_JSON_SCHEMA },
          },
          thinking: { type: "disabled" },
        } as Record<string, unknown>),
      } as unknown as Anthropic.Messages.MessageCreateParamsNonStreaming,
    })),
  });
  return batch.id;
}

export async function isBatchComplete(batchId: string): Promise<boolean> {
  const batch = await anthropic.messages.batches.retrieve(batchId);
  return batch.processing_status === "ended";
}

/**
 * Ingest batch results: store the raw response and the parsed per-dimension
 * scores (full audit trail, CLAUDE.md §6). Idempotent via the unique
 * (taskId, dimension, model, promptVersion, sampleIndex) constraint.
 */
export async function ingestBatchResults(
  batchId: string,
  model: string,
  jobIdToTaskId: ReadonlyMap<string, string>,
): Promise<{ ingested: number; failed: number }> {
  let ingested = 0;
  let failed = 0;

  for await (const entry of await anthropic.messages.batches.results(batchId)) {
    const taskId = jobIdToTaskId.get(entry.custom_id);
    if (!taskId || entry.result.type !== "succeeded") {
      failed += 1;
      continue;
    }
    const message = entry.result.message;
    // A max_tokens stop truncates the JSON mid-string — count as failed; the
    // task stays unrated and is resubmitted by the next rating run.
    if (message.stop_reason === "max_tokens") {
      failed += 1;
      continue;
    }
    const textBlock = message.content.find((block) => block.type === "text");
    if (!textBlock || textBlock.type !== "text") {
      failed += 1;
      continue;
    }

    let json: unknown;
    try {
      json = JSON.parse(textBlock.text);
    } catch {
      failed += 1;
      continue;
    }
    const parsed = ratingResponseSchema.safeParse(json);
    if (!parsed.success) {
      failed += 1;
      continue;
    }

    for (const [dimension, value] of Object.entries(parsed.data)) {
      await prisma.taskRating.upsert({
        where: {
          taskId_dimension_model_promptVersion_sampleIndex: {
            taskId,
            dimension,
            model,
            promptVersion: RATER_PROMPT_VERSION,
            sampleIndex: 0,
          },
        },
        update: {},
        create: {
          taskId,
          dimension,
          model,
          promptVersion: RATER_PROMPT_VERSION,
          sampleIndex: 0,
          rating: value.rating,
          rationale: value.rationale,
          rawResponse: JSON.parse(JSON.stringify(message)),
        },
      });
      ingested += 1;
    }
  }

  return { ingested, failed };
}