OmniRoute/src/lib/db/batches.ts

import { getDbInstance, rowToCamel, objToSnake } from "./core";
import { deleteFile } from "./files";
import { v4 as uuidv4 } from "uuid";

function parseBatchRow(row: any): BatchRecord {
  const camel = rowToCamel(row) as any;
  if (camel.metadata && typeof camel.metadata === "string") {
    try {
      camel.metadata = JSON.parse(camel.metadata);
    } catch {
      camel.metadata = null;
    }
  }
  if (camel.errors && typeof camel.errors === "string") {
    try {
      camel.errors = JSON.parse(camel.errors);
    } catch {
      camel.errors = null;
    }
  }
  if (camel.usage && typeof camel.usage === "string") {
    try {
      camel.usage = JSON.parse(camel.usage);
    } catch {
      camel.usage = null;
    }
  }
  // Normalize numeric date fields to ensure they are valid numbers
  const coerceNum = (v: any): number | null => {
    if (typeof v === "number" && Number.isFinite(v)) return v;
    if (v == null) return null;
    const n = Number(v);
    return Number.isFinite(n) ? n : null;
  };

  camel.createdAt = coerceNum(camel.createdAt) ?? 0;
  camel.inProgressAt = coerceNum(camel.inProgressAt);
  camel.expiresAt = coerceNum(camel.expiresAt);
  camel.finalizingAt = coerceNum(camel.finalizingAt);
  camel.completedAt = coerceNum(camel.completedAt);
  camel.failedAt = coerceNum(camel.failedAt);
  camel.expiredAt = coerceNum(camel.expiredAt);
  camel.cancellingAt = coerceNum(camel.cancellingAt);
  camel.cancelledAt = coerceNum(camel.cancelledAt);
  return camel as BatchRecord;
}

export interface BatchRecord {
  id: string;
  endpoint: string;
  completionWindow: string;
  status:
    | "validating"
    | "failed"
    | "in_progress"
    | "finalizing"
    | "completed"
    | "expired"
    | "cancelling"
    | "cancelled";
  inputFileId: string;
  outputFileId?: string | null;
  errorFileId?: string | null;
  createdAt: number;
  inProgressAt?: number | null;
  expiresAt?: number | null;
  finalizingAt?: number | null;
  completedAt?: number | null;
  failedAt?: number | null;
  expiredAt?: number | null;
  cancellingAt?: number | null;
  cancelledAt?: number | null;
  requestCountsTotal: number;
  requestCountsCompleted: number;
  requestCountsFailed: number;
  metadata?: Record<string, any> | null;
  apiKeyId?: string | null;
  errors?: any | null;
  model?: string | null;
  usage?: any | null;
  outputExpiresAfterSeconds?: number | null;
  outputExpiresAfterAnchor?: string | null;
}

export type BatchItemCheckpointStatus = "pending" | "processing" | "completed" | "errored";

export interface BatchItemCheckpoint {
  batchId: string;
  lineNumber: number;
  customId: string | null;
  status: BatchItemCheckpointStatus;
  result: any | null;
  error: any | null;
  createdAt: number;
  updatedAt: number;
}

function parseJsonColumn(value: unknown): any | null {
  if (value == null) return null;
  if (typeof value !== "string") return value;
  try {
    return JSON.parse(value);
  } catch {
    return null;
  }
}

function parseBatchItemCheckpoint(row: any): BatchItemCheckpoint {
  return {
    batchId: row.batch_id,
    lineNumber: Number(row.line_number),
    customId: row.custom_id ?? null,
    status: row.status,
    result: parseJsonColumn(row.result_json),
    error: parseJsonColumn(row.error_json),
    createdAt: Number(row.created_at),
    updatedAt: Number(row.updated_at),
  };
}

export function createBatch(
  batch: Omit<
    BatchRecord,
    | "id"
    | "createdAt"
    | "requestCountsTotal"
    | "requestCountsCompleted"
    | "requestCountsFailed"
    | "status"
  > & { status?: BatchRecord["status"] }
): BatchRecord {
  const db = getDbInstance();
  const id = "batch_" + uuidv4().replaceAll("-", "").substring(0, 24);
  const createdAt = Math.floor(Date.now() / 1000);
  const record: BatchRecord = {
    ...batch,
    id,
    createdAt,
    status: batch.status || "validating",
    requestCountsTotal: 0,
    requestCountsCompleted: 0,
    requestCountsFailed: 0,
    errors: batch.errors || null,
    model: batch.model || null,
    usage: batch.usage || null,
    outputExpiresAfterSeconds: batch.outputExpiresAfterSeconds || null,
    outputExpiresAfterAnchor: batch.outputExpiresAfterAnchor || null,
  };

  const snakeRecord = objToSnake({
    ...record,
    metadata: record.metadata ? JSON.stringify(record.metadata) : null,
    errors: record.errors ? JSON.stringify(record.errors) : null,
    usage: record.usage ? JSON.stringify(record.usage) : null,
  }) as any;
  const keys = Object.keys(snakeRecord);
  const values = Object.values(snakeRecord);
  const placeholders = keys.map(() => "?").join(", ");

  db.prepare(`INSERT INTO batches (${keys.join(", ")}) VALUES (${placeholders})`).run(...values);

  return record;
}

export function getBatch(id: string): BatchRecord | null {
  const db = getDbInstance();
  const row = db.prepare("SELECT * FROM batches WHERE id = ?").get(id);
  if (!row) return null;
  return parseBatchRow(row);
}

export function updateBatch(id: string, updates: Partial<BatchRecord>): boolean {
  const db = getDbInstance();
  const snakeUpdates = objToSnake(updates) as any;
  if (snakeUpdates.metadata && typeof snakeUpdates.metadata !== "string") {
    snakeUpdates.metadata = JSON.stringify(snakeUpdates.metadata);
  }
  if (snakeUpdates.errors && typeof snakeUpdates.errors !== "string") {
    snakeUpdates.errors = JSON.stringify(snakeUpdates.errors);
  }
  if (snakeUpdates.usage && typeof snakeUpdates.usage !== "string") {
    snakeUpdates.usage = JSON.stringify(snakeUpdates.usage);
  }

  const keys = Object.keys(snakeUpdates);
  if (keys.length === 0) return false;

  const setClause = keys.map((k) => `${k} = ?`).join(", ");
  const values = Object.values(snakeUpdates);

  const result = db.prepare(`UPDATE batches SET ${setClause} WHERE id = ?`).run(...values, id);
  return result.changes > 0;
}

export function ensureBatchItemCheckpoints(
  batchId: string,
  items: Array<{ lineNumber: number; customId: string | null }>
): void {
  if (items.length === 0) return;

  const db = getDbInstance();
  const now = Math.floor(Date.now() / 1000);
  const insert = db.prepare(`
    INSERT OR IGNORE INTO batch_item_checkpoints (
      batch_id,
      line_number,
      custom_id,
      status,
      result_json,
      error_json,
      created_at,
      updated_at
    )
    VALUES (?, ?, ?, 'pending', NULL, NULL, ?, ?)
  `);

  const tx = db.transaction(() => {
    for (const item of items) {
      insert.run(batchId, item.lineNumber, item.customId, now, now);
    }
  });
  tx();
}

export function countBatchItemCheckpoints(batchId: string): number {
  const db = getDbInstance();
  const row = db
    .prepare("SELECT COUNT(*) AS c FROM batch_item_checkpoints WHERE batch_id = ?")
    .get(batchId) as { c: number } | undefined;
  return row ? Number(row.c) : 0;
}

export function listBatchItemCheckpoints(batchId: string): BatchItemCheckpoint[] {
  const db = getDbInstance();
  const rows = db
    .prepare(
      `
      SELECT batch_id, line_number, custom_id, status, result_json, error_json, created_at, updated_at
      FROM batch_item_checkpoints
      WHERE batch_id = ?
      ORDER BY line_number ASC
    `
    )
    .all(batchId);
  return rows.map((row) => parseBatchItemCheckpoint(row));
}

export function markBatchItemProcessing(
  batchId: string,
  item: { lineNumber: number; customId: string | null }
): void {
  const db = getDbInstance();
  const now = Math.floor(Date.now() / 1000);
  db.prepare(
    `
    INSERT INTO batch_item_checkpoints (
      batch_id,
      line_number,
      custom_id,
      status,
      result_json,
      error_json,
      created_at,
      updated_at
    )
    VALUES (?, ?, ?, 'processing', NULL, NULL, ?, ?)
    ON CONFLICT(batch_id, line_number) DO UPDATE SET
      custom_id = excluded.custom_id,
      status = 'processing',
      result_json = NULL,
      error_json = NULL,
      updated_at = excluded.updated_at
  `
  ).run(batchId, item.lineNumber, item.customId, now, now);
}

export function markBatchItemResult(
  batchId: string,
  item: { lineNumber: number; customId: string | null },
  result: any
): void {
  const db = getDbInstance();
  const now = Math.floor(Date.now() / 1000);
  db.prepare(
    `
    UPDATE batch_item_checkpoints
    SET custom_id = ?,
        status = 'completed',
        result_json = ?,
        error_json = NULL,
        updated_at = ?
    WHERE batch_id = ? AND line_number = ?
  `
  ).run(item.customId, JSON.stringify(result), now, batchId, item.lineNumber);
}

export function markBatchItemError(
  batchId: string,
  item: { lineNumber: number; customId: string | null },
  error: any
): void {
  const db = getDbInstance();
  const now = Math.floor(Date.now() / 1000);
  db.prepare(
    `
    UPDATE batch_item_checkpoints
    SET custom_id = ?,
        status = 'errored',
        result_json = NULL,
        error_json = ?,
        updated_at = ?
    WHERE batch_id = ? AND line_number = ?
  `
  ).run(item.customId, JSON.stringify(error), now, batchId, item.lineNumber);
}

export function listBatches(apiKeyId?: string, limit: number = 20, after?: string): BatchRecord[] {
  const db = getDbInstance();
  const afterBatch = after ? getBatch(after) : null;
  let rows: any[];
  if (apiKeyId) {
    if (afterBatch) {
      rows = db
        .prepare(
          "SELECT * FROM batches WHERE api_key_id = ? AND (created_at < ? OR (created_at = ? AND id < ?)) ORDER BY created_at DESC, id DESC LIMIT ?"
        )
        .all(apiKeyId, afterBatch.createdAt, afterBatch.createdAt, after, limit);
    } else {
      rows = db
        .prepare(
          "SELECT * FROM batches WHERE api_key_id = ? ORDER BY created_at DESC, id DESC LIMIT ?"
        )
        .all(apiKeyId, limit);
    }
  } else if (afterBatch) {
    rows = db
      .prepare(
        "SELECT * FROM batches WHERE (created_at < ? OR (created_at = ? AND id < ?)) ORDER BY created_at DESC, id DESC LIMIT ?"
      )
      .all(afterBatch.createdAt, afterBatch.createdAt, after, limit);
  } else {
    rows = db.prepare("SELECT * FROM batches ORDER BY created_at DESC, id DESC LIMIT ?").all(limit);
  }
  return rows.map((row) => parseBatchRow(row));
}

export function countBatches(apiKeyId?: string): number {
  const db = getDbInstance();
  if (apiKeyId) {
    const row = db
      .prepare("SELECT COUNT(*) as c FROM batches WHERE api_key_id = ?")
      .get(apiKeyId) as { c: number } | undefined;
    return row ? Number(row.c) : 0;
  } else {
    const row = db.prepare("SELECT COUNT(*) as c FROM batches").get() as { c: number } | undefined;
    return row ? Number(row.c) : 0;
  }
}

export function getPendingBatches(): BatchRecord[] {
  const db = getDbInstance();
  const rows = db
    .prepare(
      "SELECT * FROM batches WHERE status IN ('validating', 'in_progress', 'finalizing', 'cancelling')"
    )
    .all();
  return rows.map((row) => parseBatchRow(row));
}

export function getTerminalBatches(): BatchRecord[] {
  const db = getDbInstance();
  const rows = db
    .prepare(
      "SELECT * FROM batches WHERE status IN ('completed', 'failed', 'cancelled', 'expired') ORDER BY created_at ASC"
    )
    .all();
  return rows.map((row) => parseBatchRow(row));
}

export function deleteBatch(id: string): boolean {
  const db = getDbInstance();
  const batch = getBatch(id);
  if (!batch) return false;

  db.prepare("DELETE FROM batch_item_checkpoints WHERE batch_id = ?").run(id);

  // Soft-delete associated files (input, output, error)
  if (batch.inputFileId) {
    try {
      deleteFile(batch.inputFileId);
    } catch {
      /* ignore */
    }
  }
  if (batch.outputFileId) {
    try {
      deleteFile(batch.outputFileId);
    } catch {
      /* ignore */
    }
  }
  if (batch.errorFileId) {
    try {
      deleteFile(batch.errorFileId);
    } catch {
      /* ignore */
    }
  }

  const result = db.prepare("DELETE FROM batches WHERE id = ?").run(id);
  return result.changes > 0;
}

/**
 * Bulk-delete completed batches and the files they reference.
 *
 * `apiKeyId` scopes EVERY statement to that owner. Omitting it keeps the
 * instance-wide sweep, which is legitimate for the operator's own dashboard
 * (session auth) and for nothing else: without the predicate, an ordinary
 * inference key could wipe every tenant's completed batches and null out their
 * file contents (GHSA-wvxc-jp3v-5mg5). Same ownership shape as `listBatches`
 * and `countBatches` above.
 */
export function deleteCompletedBatches(apiKeyId?: string | null): {
  deletedBatches: number;
  deletedFiles: number;
} {
  const db = getDbInstance();
  const scoped = typeof apiKeyId === "string" && apiKeyId.length > 0;

  // Collect unique file IDs from the completed batches in scope
  const rows = (
    scoped
      ? db
          .prepare(
            "SELECT input_file_id, output_file_id, error_file_id FROM batches WHERE status = 'completed' AND api_key_id = ?"
          )
          .all(apiKeyId)
      : db
          .prepare(
            "SELECT input_file_id, output_file_id, error_file_id FROM batches WHERE status = 'completed'"
          )
          .all()
  ) as Array<{
    input_file_id: string | null;
    output_file_id: string | null;
    error_file_id: string | null;
  }>;

  const fileIds = new Set<string>();
  for (const row of rows) {
    if (row.input_file_id) fileIds.add(row.input_file_id);
    if (row.output_file_id) fileIds.add(row.output_file_id);
    if (row.error_file_id) fileIds.add(row.error_file_id);
  }

  let deletedFiles = 0;
  for (const fid of fileIds) {
    try {
      if (deleteFile(fid)) deletedFiles++;
    } catch {
      /* ignore */
    }
  }

  if (scoped) {
    db.prepare(
      "DELETE FROM batch_item_checkpoints WHERE batch_id IN (SELECT id FROM batches WHERE status = 'completed' AND api_key_id = ?)"
    ).run(apiKeyId);
    const result = db
      .prepare("DELETE FROM batches WHERE status = 'completed' AND api_key_id = ?")
      .run(apiKeyId);
    return { deletedBatches: result.changes, deletedFiles };
  }

  db.prepare(
    "DELETE FROM batch_item_checkpoints WHERE batch_id IN (SELECT id FROM batches WHERE status = 'completed')"
  ).run();

  const result = db.prepare("DELETE FROM batches WHERE status = 'completed'").run();
  return { deletedBatches: result.changes, deletedFiles };
}