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 function createBatch(
  batch: Omit<
    BatchRecord,
    "id" | "createdAt" | "requestCountsTotal" | "requestCountsCompleted" | "requestCountsFailed"
  >
): 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 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;

  // 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;
}

export function deleteCompletedBatches(): { deletedBatches: number; deletedFiles: number } {
  const db = getDbInstance();

  // Collect unique file IDs from all completed batches
  const rows = 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 */
    }
  }

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