OmniRoute/open-sse/services/batchProcessor.ts

import { v4 as uuidv4 } from "uuid";
import type { BatchItemCheckpoint, BatchRecord } from "@/lib/localDb";
import {
  countBatchItemCheckpoints,
  createFile,
  deleteFile,
  ensureBatchItemCheckpoints,
  getApiKeyById,
  getBatch,
  getFileContent,
  getPendingBatches,
  getTerminalBatches,
  listBatchItemCheckpoints,
  listFiles,
  markBatchItemError,
  markBatchItemProcessing,
  markBatchItemResult,
  updateBatch,
} from "@/lib/localDb";
import { dispatch } from "@/lib/batches/dispatch";
import type { SupportedBatchEndpoint } from "@/shared/constants/batchEndpoints";
import { DEFAULT_BATCH_EXPIRATION_SECONDS } from "@/shared/constants/batch";

let isProcessing: boolean = false;
let pollInterval: NodeJS.Timeout | null = null;
const activeProcesses = new Set<Promise<void>>();
const activeBatches = new Set<string>();
const DEFAULT_BATCH_WINDOW_SECONDS: number = 24 * 60 * 60;
const BATCH_RETRY_DURATION_MS: number =
  Number.parseInt(process.env.BATCH_RETRY_DURATION_MS ?? "", 10) || 24 * 60 * 60 * 1_000;
const BATCH_BACKOFF_BASE_MS: number =
  Number.parseInt(process.env.BATCH_BACKOFF_BASE_MS ?? "", 10) || 5_000;
const BATCH_BACKOFF_MAX_MS: number =
  Number.parseInt(process.env.BATCH_BACKOFF_MAX_MS ?? "", 10) || 3_600_000;
const BATCH_MAX_CONCURRENT: number =
  Number.parseInt(process.env.BATCH_MAX_CONCURRENT ?? "", 10) || 1;

interface BatchRequestItem {
  body: Record<string, unknown>;
  customId: string | null;
  lineNumber: number;
  method: "POST";
  url: SupportedBatchEndpoint;
}

export function initBatchProcessor() {
  if (pollInterval) return pollInterval;
  console.log("[BATCH] Initializing batch processor polling...");

  pollInterval = setInterval(async (): Promise<void> => {
    if (isProcessing) return;
    try {
      isProcessing = true;
      await processPendingBatches();
    } catch (err) {
      console.error("[BATCH] Polling error:", err);
    } finally {
      isProcessing = false;
    }
  }, 10000).unref(); // Poll every 10s; unref'd so it never keeps the process alive
  return pollInterval;
}

export function stopBatchProcessor(): void {
  if (pollInterval) {
    clearInterval(pollInterval);
    pollInterval = null;
    console.log("[BATCH] Stopped batch processor polling.");
  }
}

export async function processPendingBatches(): Promise<void> {
  const pending = getPendingBatches();

  // Phase 1: Stale recovery — in_progress/finalizing batches not in activeBatches
  // are from a previous session; reset checkpointed batches to validating so they
  // can be completed without replaying already-dispatched items.
  for (const batch of pending) {
    if (batch.status === "in_progress" || batch.status === "finalizing") {
      if (!activeBatches.has(batch.id)) {
        recoverStaleBatch(batch);
      }
    }
  }

  // Phase 2: Process actions respecting concurrency limit
  const remaining = getPendingBatches(); // re-fetch after recovery updates
  let activeCount = activeBatches.size;

  for (const batch of remaining) {
    if (batch.status === "validating") {
      if (activeCount >= BATCH_MAX_CONCURRENT) {
        console.log(
          `[BATCH] Concurrency limit ${BATCH_MAX_CONCURRENT} reached, deferring batch ${batch.id}`
        );
        continue;
      }
      activeCount++;
      await startBatch(batch);
    } else if (batch.status === "cancelling") {
      await cancelBatch(batch);
    }
  }

  // Cleanup task: delete files for batches completed more than completionWindow ago
  await cleanupExpiredBatches();
}

function recoverStaleBatch(batch: BatchRecord): void {
  const checkpointCount = countBatchItemCheckpoints(batch.id);
  const hasPotentialExternalEffects =
    batch.requestCountsTotal > 0 ||
    batch.requestCountsCompleted > 0 ||
    batch.requestCountsFailed > 0 ||
    batch.status === "finalizing";

  if (checkpointCount === 0 && hasPotentialExternalEffects) {
    console.warn(
      `[BATCH] Stale batch ${batch.id} has no item checkpoints; failing instead of replaying provider calls`
    );
    failBatch(
      batch.id,
      "Cannot safely recover stale batch because item checkpoints are unavailable; create a new batch to retry intentionally."
    );
    return;
  }

  console.log(`[BATCH] Recovering stale batch ${batch.id} (${batch.status}) → validating`);

  if (batch.outputFileId) {
    deleteFile(batch.outputFileId);
  }
  if (batch.errorFileId) {
    deleteFile(batch.errorFileId);
  }

  updateBatch(batch.id, {
    status: "validating",
    inProgressAt: null,
    finalizingAt: null,
    outputFileId: null,
    errorFileId: null,
    ...(checkpointCount === 0
      ? {
          requestCountsCompleted: 0,
          requestCountsFailed: 0,
        }
      : {}),
  });
}

function parseBatchWindowSeconds(window: string | null | undefined): number {
  if (!window) return DEFAULT_BATCH_WINDOW_SECONDS;
  const match = /^(\d+)([hdm])$/.exec(window);
  if (!match) return DEFAULT_BATCH_WINDOW_SECONDS;

  const value = Number.parseInt(match[1], 10);
  const unit = match[2];
  if (unit === "h") return value * 3600;
  if (unit === "d") return value * 86400;
  if (unit === "m") return value * 60;
  return DEFAULT_BATCH_WINDOW_SECONDS;
}

function getBatchOutputExpiresAt(batch: BatchRecord): number | null {
  if (
    batch.outputExpiresAfterAnchor === "created_at" &&
    typeof batch.outputExpiresAfterSeconds === "number" &&
    batch.outputExpiresAfterSeconds > 0
  ) {
    return batch.createdAt + batch.outputExpiresAfterSeconds;
  }

  const completionTime: number =
    batch.completedAt || batch.failedAt || batch.cancelledAt || batch.expiredAt;
  if (!completionTime) return null;
  // Default: batch output files expire 30 days after completion
  return completionTime + DEFAULT_BATCH_EXPIRATION_SECONDS;
}

function resolveBatchApiKeyValue(batch: Pick<BatchRecord, "apiKeyId">, apiKeyRow: any): any {
  if (typeof apiKeyRow?.key === "string" && apiKeyRow.key.length > 0) {
    return apiKeyRow.key;
  }
  if (batch.apiKeyId === "env-key") {
    return process.env.OMNIROUTE_API_KEY || process.env.ROUTER_API_KEY || null;
  }
  return null;
}

export function parseBatchItems(
  content: Buffer,
  batchEndpoint: SupportedBatchEndpoint
): { items: BatchRequestItem[]; error: null } | { items: null; error: string } {
  const lines = content
    .toString()
    .split("\n")
    .map((line) => line.trim())
    .filter(Boolean);

  const items: BatchRequestItem[] = [];
  for (const [index, line] of lines.entries()) {
    let parsed: Record<string, unknown>;
    try {
      parsed = JSON.parse(line);
    } catch {
      return { items: null, error: `Line ${index + 1} is not valid JSON` };
    }

    const method = String(parsed.method || "POST").toUpperCase();
    const url = parsed.url;
    const body = parsed.body;

    if (method !== "POST") {
      return {
        items: null,
        error: `Line ${index + 1} uses unsupported method ${method}; only POST is supported`,
      };
    }
    if (url !== batchEndpoint) {
      return {
        items: null,
        error: `Line ${index + 1} url ${String(url)} does not match batch endpoint ${batchEndpoint}`,
      };
    }
    if (!body || typeof body !== "object" || Array.isArray(body)) {
      return { items: null, error: `Line ${index + 1} must include a JSON object body` };
    }

    items.push({
      body: body as Record<string, unknown>,
      customId: typeof parsed.custom_id === "string" ? parsed.custom_id : null,
      lineNumber: index + 1,
      method: "POST",
      url: batchEndpoint,
    });
  }

  return { items, error: null };
}

async function cleanupExpiredBatches(): Promise<void> {
  try {
    const now = Math.floor(Date.now() / 1_000);
    const batches = getTerminalBatches();

    // Delete files for terminal batches that have exceeded their completion window
    for (const batch of batches) {
      const completionTime =
        batch.completedAt || batch.failedAt || batch.cancelledAt || batch.expiredAt;
      // Input files expire 30 days after batch completion
      const inputExpiresAt =
        completionTime && batch.inputFileId
          ? completionTime + DEFAULT_BATCH_EXPIRATION_SECONDS
          : null;
      const outputExpiresAt = getBatchOutputExpiresAt(batch);

      if (batch.inputFileId && inputExpiresAt && now > inputExpiresAt) {
        deleteFile(batch.inputFileId);
      }
      if (batch.outputFileId && outputExpiresAt && now > outputExpiresAt) {
        deleteFile(batch.outputFileId);
      }
      if (batch.errorFileId && outputExpiresAt && now > outputExpiresAt) {
        deleteFile(batch.errorFileId);
      }
    }

    // Expire validating batches that have exceeded their completion window
    for (const batch of getPendingBatches()) {
      if (batch.status === "validating") {
        const windowSeconds = parseBatchWindowSeconds(batch.completionWindow);
        if (now - batch.createdAt > windowSeconds) {
          updateBatch(batch.id, { status: "expired", expiredAt: now });
        }
      }
    }

    // Cleanup orphan files (batch-purpose files stuck in validating after 48h)
    // Use asc order so oldest files are processed first; use a high limit to avoid missing old orphans.
    const allFiles = listFiles({ order: "asc", limit: 100 });
    for (const file of allFiles) {
      if (file.purpose === "batch" && now - file.createdAt > DEFAULT_BATCH_EXPIRATION_SECONDS) {
        deleteFile(file.id);
      }
    }
  } catch (err) {
    console.error("[BATCH] Cleanup error:", err);
  }
}

async function startBatch(batch: any): Promise<void> {
  console.log(`[BATCH] Starting batch ${batch.id}`);

  const content = getFileContent(batch.inputFileId);
  if (!content) {
    failBatch(batch.id, "Input file content not found");
    return;
  }

  try {
    const parsedItems = parseBatchItems(content, batch.endpoint);
    if (parsedItems.error) {
      // Set total count even on validation failure so UI shows correct numbers
      const lines = content
        .toString()
        .split("\n")
        .map((line) => line.trim())
        .filter(Boolean);
      updateBatch(batch.id, {
        requestCountsTotal: lines.length,
        requestCountsFailed: lines.length, // All failed due to validation error
      });
      failBatch(batch.id, parsedItems.error);
      return;
    }
    const total = parsedItems.items.length;

    console.log(`[BATCH] Batch ${batch.id} contains (${total} items)`);

    ensureBatchItemCheckpoints(batch.id, parsedItems.items);

    updateBatch(batch.id, {
      status: "in_progress",
      inProgressAt: Math.floor(Date.now() / 1000),
      requestCountsTotal: total,
    });

    activeBatches.add(batch.id);

    // Fire-and-forget: process items in the background so the poll loop isn't blocked.
    // isProcessing prevents a second poll tick from overlapping.
    const p = processBatchItems(batch, parsedItems.items).catch((err) => {
      console.error(`[BATCH] Critical error in processBatchItems for ${batch.id}:`, err);
      failBatch(batch.id, String(err));
    });
    activeProcesses.add(p);
    p.finally(() => {
      activeProcesses.delete(p);
      activeBatches.delete(batch.id);
    });
  } catch (err) {
    console.error(`[BATCH] Error starting batch ${batch.id}:`, err);
    failBatch(batch.id, err instanceof Error ? err.message : String(err));
  }
}

let prevHeaders: Headers | null = null;
let prevHeadersTimestamp: number = 0;
const HEADERS_CACHE_TTL_MS = 60_000;

async function processBatchItems(batch: BatchRecord, items: BatchRequestItem[]): Promise<void> {
  const state = createBatchState(batch);
  const checkpoints = new Map<number, BatchItemCheckpoint>(
    listBatchItemCheckpoints(batch.id).map((checkpoint) => [checkpoint.lineNumber, checkpoint])
  );

  const apiKey = await resolveApiKey(batch);

  for (const item of items) {
    if (isBatchCancelled(batch.id)) break;

    const checkpoint = checkpoints.get(item.lineNumber);
    if (checkpoint && applyRecoveredCheckpoint(batch.id, item, checkpoint, state)) {
      maybePersistProgress(batch.id, state);
      continue;
    }

    const cachedHeaders =
      prevHeaders && Date.now() - prevHeadersTimestamp < HEADERS_CACHE_TTL_MS ? prevHeaders : null;
    if (cachedHeaders) {
      const delay = maybeThrottle(cachedHeaders);
      if (delay) {
        await sleep(delay);
      }
    }

    markBatchItemProcessing(batch.id, item);

    try {
      const response = await processSingleItemWithRetry(item, apiKey);
      let responseBody: unknown;
      try {
        responseBody = await response.clone().json();
      } catch {
        responseBody = await response.text();
      }

      // Record the item's result so finalization can emit output/error files.
      // The output file format expects entries like:
      // { id, custom_id, response: { status_code, body } }
      const wrapped = {
        id: `req_${uuidv4().replaceAll("-", "")}`,
        custom_id: item.customId ?? null,
        response: {
          status_code: response.status,
          body: responseBody,
        },
      };

      markBatchItemResult(batch.id, item, wrapped);
      state.results.push(wrapped);
      applyItemResult(state, response.status, responseBody);
      prevHeaders = response.headers;
      prevHeadersTimestamp = Date.now();
    } catch (exception) {
      // Track processing-level errors separately (items that failed to be processed)
      const error = { custom_id: item.customId ?? null, error: String(exception) };
      markBatchItemError(batch.id, item, error);
      state.errors.push(error);
      state.failed++;
      prevHeaders = null;
      prevHeadersTimestamp = 0;
    }

    maybePersistProgress(batch.id, state);
  }

  return finalizeBatch(batch.id, state.results, state.errors);
}

function applyRecoveredCheckpoint(
  batchId: string,
  item: BatchRequestItem,
  checkpoint: BatchItemCheckpoint,
  state: ReturnType<typeof createBatchState>
): boolean {
  if (checkpoint.status === "completed" && checkpoint.result) {
    state.results.push(checkpoint.result);
    applyItemResult(
      state,
      checkpoint.result.response?.status_code ?? 500,
      checkpoint.result.response?.body
    );
    return true;
  }

  if (checkpoint.status === "errored" && checkpoint.error) {
    state.errors.push(checkpoint.error);
    state.failed++;
    return true;
  }

  if (checkpoint.status === "processing") {
    const error = {
      custom_id: item.customId ?? null,
      error:
        "Batch item was interrupted before its provider response was recorded; it was not replayed to avoid duplicate provider work.",
    };
    markBatchItemError(batchId, item, error);
    state.errors.push(error);
    state.failed++;
    return true;
  }

  return false;
}

function isBatchCancelled(batchId: string): boolean {
  const current = getBatch(batchId);

  return !current || current.status === "cancelling" || current.status === "cancelled";
}

async function resolveApiKey(batch: BatchRecord): Promise<any> {
  const apiKeyRow = batch.apiKeyId ? await getApiKeyById(batch.apiKeyId) : null;
  return resolveBatchApiKeyValue(batch, apiKeyRow);
}

async function processSingleItemWithRetry(item: BatchRequestItem, apiKey: string) {
  // Time-based retry limit: individual batch items can retry for up to 24 hours.
  // This accommodates large batches against heavily rate-limited providers.
  const MAX_RETRY_DURATION_MS = BATCH_RETRY_DURATION_MS;
  const maxRetries = 200; // safety ceiling — time limit should kick in first
  const retryStartedAt = Date.now();

  let response: Response = null;
  for (let attempt = 1; attempt <= maxRetries; attempt++) {
    // If the previous attempt got a response, check for rate limit headers and throttle if needed before retrying.
    if (response) {
      const delay = maybeThrottle(response.headers);
      if (delay) {
        await sleep(delay);
      }
    }

    response = await processSingleItem(item, apiKey);

    if (
      (response.status === 429 || response.status === 502 || response.status === 504) &&
      attempt < maxRetries
    ) {
      // Bail if we've been retrying for longer than the allowed window
      if (Date.now() - retryStartedAt >= MAX_RETRY_DURATION_MS) {
        console.warn(
          `[BATCH] Item ${item.customId ?? "(no id)"} exceeded 24h retry window after ${attempt} attempts — giving up`
        );
        return response;
      }
      const delay = getRetryDelayMs(response.headers) ?? getBackoffDelayMs(attempt);
      await sleep(delay);
      continue;
    }

    return response;
  }
}

async function processSingleItem(item: BatchRequestItem, apiKey: string) {
  const body = buildRequestBody(item);

  return await dispatch.dispatchBatchApiRequest({
    endpoint: item.url,
    body,
    apiKey,
  });
}

export function buildRequestBody(item: BatchRequestItem) {
  const isChatEndpoint = ![
    "/v1/embeddings",
    "/v1/moderations",
    "/v1/images/generations",
    "/v1/images/edits",
    "/v1/videos",
    "/v1/videos/generations",
  ].includes(item.url);

  return {
    ...item.body,
    ...(isChatEndpoint ? { stream: false } : {}),
  };
}

function getBackoffDelayMs(attempt: number) {
  const baseMs = BATCH_BACKOFF_BASE_MS;
  const maxMs = BATCH_BACKOFF_MAX_MS;

  // exponential: 2^attempt * base
  const exp = Math.min(maxMs, baseMs * 2 ** attempt);

  // jitter ±20%
  const jitterFactor = 1 + (Math.random() * 0.4 - 0.2);

  return Math.floor(exp * jitterFactor);
}

function sleep(ms: number) {
  return new Promise((res) => setTimeout(res, ms));
}
function getRetryDelayMs(headers: Headers): number | null {
  const retryAfter = headers.get("retry-after");
  if (retryAfter) {
    const seconds = Number(retryAfter);
    if (!Number.isNaN(seconds)) {
      return seconds * 1_000;
    }

    // fallback: HTTP-date
    const date = new Date(retryAfter).getTime();
    if (!Number.isNaN(date)) {
      return Math.max(0, date - Date.now());
    }
  }

  return null;
}

export function maybeThrottle(headers: Headers): number | null {
  // Mistral reports these headers from their API
  const remainingReq = toNumber(headers.get("x-ratelimit-remaining-req-minute"));
  const limitReq = toNumber(headers.get("x-ratelimit-limit-req-minute"));

  const remainingTokens = toNumber(headers.get("x-ratelimit-remaining-tokens-minute"));
  const cost = toNumber(headers.get("x-ratelimit-tokens-query-cost"));

  let pressures: number[] = [];

  // Request pressureRemaining
  if (remainingReq !== null && limitReq !== null) {
    if (limitReq > 0) {
      pressures.push(remainingReq / limitReq);
    }
  }

  // Token pressureRemaining
  if (remainingTokens !== null && cost !== null) {
    if (remainingTokens + cost > 0) {
      pressures.push(remainingTokens / (remainingTokens + cost));
    }
  }

  if (pressures.length === 0) {
    console.log("[BATCH] Throttle check - no rate-limit headers present");
    return null;
  } else {
    const tokenTotal = remainingTokens != null && cost != null ? remainingTokens + cost : null;
    console.log(
      `[BATCH] Throttle check - Request pressure: ${remainingReq ?? "n/a"}/${limitReq ?? "n/a"}, Token pressure: ${remainingTokens ?? "n/a"}/${tokenTotal ?? "n/a"}`
    );
  }

  const pressureRemaining = Math.min(...pressures);

  const delay = throttleDelay(pressureRemaining);
  if (delay !== null) {
    console.log(
      `[BATCH] Throttling next request with delay of ${Math.round(delay)}ms (pressure remaining: ${(pressureRemaining * 100).toFixed(2)}%)`
    );
  }
  return delay;
}

function throttleDelay(pressure: number): number | null {
  if (pressure >= 0.2) return null;

  const severity = (0.2 - pressure) / 0.2;

  const delay = Math.pow(severity, 2) * 30_000;

  return 200 + delay + Math.random() * 1000;
}

const toNumber = (v: string | null) => {
  if (v === null || v.trim() === "") return null;
  const n = Number(v);
  return Number.isFinite(n) ? n : null;
};

function createBatchState(batch: BatchRecord) {
  return {
    results: [],
    errors: [],
    completed: 0,
    failed: 0,
    tokens: {
      input: 0,
      output: 0,
      reasoning: 0,
    },
    model: batch.model || null,
  };
}

function applyItemResult(state: any, statusCode: number, body: any): void {
  if (statusCode >= 400 || body?.error) {
    state.failed++;
  } else {
    state.completed++;

    if (body?.usage) {
      state.tokens.input += body.usage.prompt_tokens || body.usage.input_tokens || 0;
      state.tokens.output += body.usage.completion_tokens || body.usage.output_tokens || 0;
      state.tokens.reasoning += body.usage.completion_tokens_details?.reasoning_tokens || 0;
    }

    if (!state.model && body?.model) {
      state.model = body.model;
    }
  }
}

function maybePersistProgress(batchId: string, state: any): void {
  // Persist basic progress (completed/failed counts + model) on every item so
  // the UI can show up-to-date progress even for small batches.
  try {
    updateBatch(batchId, {
      requestCountsCompleted: state.completed,
      requestCountsFailed: state.failed,
      model: state.model,
    });
  } catch (err) {
    console.error(`[BATCH] Failed to persist progress for ${batchId}:`, err);
  }

  // Persist richer usage/model info less frequently to avoid excessive writes.
  const total = state.completed + state.failed;
  if (total % 50 !== 0) return;

  try {
    updateBatch(batchId, {
      requestCountsCompleted: state.completed,
      requestCountsFailed: state.failed,
      model: state.model,
      usage: {
        input_tokens: state.tokens.input,
        output_tokens: state.tokens.output,
        total_tokens: state.tokens.input + state.tokens.output,
        input_tokens_details: { cached_tokens: 0 },
        output_tokens_details: { reasoning_tokens: state.tokens.reasoning },
      },
    });
  } catch (err) {
    console.error(`[BATCH] Failed to persist extended progress for ${batchId}:`, err);
  }
}

async function finalizeBatch(
  batchId: string,
  results: any[],
  itemsWithErrors: any[]
): Promise<void> {
  const current = getBatch(batchId);

  if (handleCancellation(batchId, current)) return;

  // Mark as finalizing first
  markFinalizing(batchId);

  // Compute counts from results
  const successes = results.filter(
    (r) =>
      typeof r.response?.status_code === "number" &&
      r.response.status_code < 400 &&
      !r.response.body?.error
  );
  const failuresFromResults = results.filter(
    (r) =>
      (typeof r.response?.status_code === "number" && r.response.status_code >= 400) ||
      r.response?.body?.error
  );
  const processingErrors = itemsWithErrors || [];

  const completedCount = successes.length;
  const failedCount = failuresFromResults.length + processingErrors.length;
  const totalCount = current?.requestCountsTotal || completedCount + failedCount;

  // Aggregate usage from per-item responses when available
  let inputTokens = 0;
  let outputTokens = 0;
  let reasoningTokens = 0;

  for (const r of results) {
    try {
      const body = r.response?.body;
      if (!body) continue;
      const usage = body.usage || {};
      inputTokens += usage.prompt_tokens ?? usage.input_tokens ?? usage.total_tokens ?? 0;
      outputTokens += usage.completion_tokens ?? 0;
      reasoningTokens += usage.completion_tokens_details?.reasoning_tokens ?? 0;
    } catch (err) {
      console.error("Failed to aggregate usage for batch", batchId, err);
    }
  }

  const model =
    results.find((r) => r.response?.body?.model)?.response?.body?.model || current?.model || null;

  const completionTime = now();

  // Persist final counts and (approximate) usage so UI shows correct numbers
  try {
    updateBatch(batchId, {
      requestCountsTotal: totalCount,
      requestCountsCompleted: completedCount,
      requestCountsFailed: failedCount,
      model,
      usage: {
        input_tokens: inputTokens,
        output_tokens: outputTokens,
        total_tokens: inputTokens + outputTokens,
        input_tokens_details: { cached_tokens: 0 },
        output_tokens_details: { reasoning_tokens: reasoningTokens },
      },
      // also set completedAt so file expiration calculation can use it
      completedAt: completionTime,
    });
  } catch (err) {
    console.error(`[BATCH] Failed to persist final progress for ${batchId}:`, err);
  }

  // Re-read the batch (with completedAt set) so file creation sees a completion timestamp
  const batchForFiles = getBatch(batchId);

  const outputFileId = createSuccessFile(batchId, batchForFiles, results);
  const errorFileId = createErrorFile(batchId, batchForFiles, results, itemsWithErrors);

  completeBatch(batchId, outputFileId, errorFileId);
}

function handleCancellation(batchId: string, current: any): boolean {
  if (!current) return true;

  if (current.status === "cancelling") {
    updateBatch(batchId, {
      status: "cancelled",
      cancelledAt: now(),
    });
    return true;
  }

  return current.status === "cancelled";
}

function markFinalizing(batchId: string): void {
  updateBatch(batchId, {
    status: "finalizing",
    finalizingAt: now(),
  });
}

function completeBatch(
  batchId: string,
  outputFileId: string | null,
  errorFileId: string | null
): void {
  updateBatch(batchId, {
    status: "completed",
    completedAt: now(),
    outputFileId,
    errorFileId,
  });

  const b = getBatch(batchId);
  const total = b?.requestCountsTotal ?? "?";
  console.log(`[BATCH] Completed batch ${batchId} (${total} items)`);
}

function now(): number {
  return Math.floor(Date.now() / 1_000);
}

function createSuccessFile(batchId: string, current: any, results: any[]): string | null {
  const successes = results.filter((r) => r.response.status_code < 400 && !r.response.body?.error);

  if (successes.length === 0) return null;

  const content = toJsonl(successes);

  const file = createFile({
    bytes: Buffer.byteLength(content),
    filename: `batch_${batchId}_output.jsonl`,
    purpose: "batch_output",
    content: Buffer.from(content),
    apiKeyId: current?.apiKeyId,
    expiresAt: getBatchOutputExpiresAt(current),
  });

  return file.id;
}

function createErrorFile(
  batchId: string,
  current: any,
  results: any[],
  itemsWithErrors: any[]
): string | null {
  const failures = results.filter((r) => r.response.status_code >= 400 || r.response.body?.error);

  const processErrors = itemsWithErrors.map((e) => ({
    id: `batch_req_${uuidv4().replaceAll("-", "")}`,
    custom_id: e.custom_id,
    response: null,
    error: { message: e.error, type: "batch_process_error" },
  }));

  const allFailures = [...failures, ...processErrors];

  if (allFailures.length === 0) return null;

  const content = toJsonl(allFailures);

  const file = createFile({
    bytes: Buffer.byteLength(content),
    filename: `batch_${batchId}_error.jsonl`,
    purpose: "batch_output",
    content: Buffer.from(content),
    apiKeyId: current?.apiKeyId,
    expiresAt: getBatchOutputExpiresAt(current),
  });

  return file.id;
}

function toJsonl(items: any[]): string {
  return items.map((i) => JSON.stringify(i)).join("\n");
}

async function cancelBatch(batch: any): Promise<void> {
  updateBatch(batch.id, {
    status: "cancelled",
    cancelledAt: Math.floor(Date.now() / 1000),
  });
  console.log(`[BATCH] Cancelled batch ${batch.id}`);
}

function failBatch(batchId: string, reason: string): void {
  updateBatch(batchId, {
    status: "failed",
    failedAt: Math.floor(Date.now() / 1000),
    errors: [{ message: reason }],
  });
  activeBatches.delete(batchId);
}

export async function waitForAllBatches(): Promise<void> {
  await Promise.all(Array.from(activeProcesses));
}

// Test helpers
export function getCachedHeaders(): { headers: Headers | null; timestamp: number } {
  return { headers: prevHeaders, timestamp: prevHeadersTimestamp };
}
export function resetCachedHeaders(): void {
  prevHeaders = null;
  prevHeadersTimestamp = 0;
}
export function resetBatchProcessorState(): void {
  activeBatches.clear();
  activeProcesses.clear();
  isProcessing = false;
  prevHeaders = null;
  prevHeadersTimestamp = 0;
}