hyb-oa/server/src/jobs/ocrJob.js

const fs = require("node:fs/promises");
const path = require("node:path");
const cron = require("node-cron");

const { runOcrForInvoice } = require("../services/ocr");

const OCR_DUPLICATE_ERROR = "duplicate_invoice_code_number";

function trimErrorMessage(error) {
  const message = error?.message ?? "ocr_failed";
  if (message.length > 500) {
    return message.slice(0, 500);
  }
  return message;
}

function toDateOnly(value) {
  if (!value) {
    return null;
  }

  return `${value}T00:00:00.000Z`;
}

function toMoney(value) {
  if (value == null) {
    return null;
  }

  const num = Number(value);
  if (!Number.isFinite(num)) {
    return null;
  }

  return num.toFixed(2);
}

async function claimPendingInvoices(prisma, config, workerId) {
  const batchSize = Math.max(1, Number(config.ocrBatchSize || 10));
  const lockTtlSeconds = Math.max(60, Number(config.ocrLockTtlSeconds || 300));

  return prisma.$queryRaw`
    WITH candidates AS (
      SELECT id
      FROM "invoice_pool"
      WHERE status = 'pending_ocr'
        AND (ocr_locked_at IS NULL OR ocr_locked_at < NOW() - make_interval(secs => ${lockTtlSeconds}))
      ORDER BY created_at ASC
      LIMIT ${batchSize}
      FOR UPDATE SKIP LOCKED
    )
    UPDATE "invoice_pool" AS i
    SET ocr_locked_at = NOW(),
        ocr_lock_owner = ${workerId}
    FROM candidates
    WHERE i.id = candidates.id
    RETURNING i.id, i.file_sha256, i.ocr_type, i.storage_path, i.ocr_attempts
  `;
}

async function findDuplicateInvoice(prisma, invoiceId, invoiceCode, invoiceNumber) {
  if (!invoiceCode || !invoiceNumber) {
    return null;
  }

  return prisma.invoicePool.findFirst({
    where: {
      id: { not: invoiceId },
      invoice_code: invoiceCode,
      invoice_number: invoiceNumber,
      is_duplicate_override: false,
      OR: [
        { status: "bound" },
        {
          status: "archived",
          archived_reason: "used_complete",
        },
      ],
    },
    select: {
      id: true,
    },
  });
}

async function getOrCreateOcrResult(prisma, config, claimedInvoice, imageBase64) {
  const where = {
    file_sha256_ocr_provider_ocr_type: {
      file_sha256: claimedInvoice.file_sha256,
      ocr_provider: "tencent",
      ocr_type: claimedInvoice.ocr_type,
    },
  };

  const cached = await prisma.ocrResult.findUnique({
    where,
    select: {
      id: true,
      result_json: true,
    },
  });

  if (cached) {
    return {
      ocrResultId: cached.id,
      normalized: cached.result_json,
    };
  }

  const normalized = await runOcrForInvoice({
    config,
    ocrType: claimedInvoice.ocr_type,
    imageBase64,
  });

  try {
    const created = await prisma.ocrResult.create({
      data: {
        file_sha256: claimedInvoice.file_sha256,
        ocr_provider: "tencent",
        ocr_type: claimedInvoice.ocr_type,
        result_json: normalized,
      },
      select: {
        id: true,
      },
    });

    return {
      ocrResultId: created.id,
      normalized,
    };
  } catch (error) {
    if (error?.code !== "P2002") {
      throw error;
    }

    const raced = await prisma.ocrResult.findUnique({
      where,
      select: {
        id: true,
        result_json: true,
      },
    });

    if (!raced) {
      throw error;
    }

    return {
      ocrResultId: raced.id,
      normalized: raced.result_json,
    };
  }
}

async function markInvoiceAsDuplicate(prisma, invoiceId, normalized, ocrResultId) {
  const now = new Date();
  await prisma.invoicePool.update({
    where: { id: invoiceId },
    data: {
      status: "ocr_failed",
      ocr_result_id: ocrResultId,
      invoice_code: normalized.invoiceCode ?? null,
      invoice_number: normalized.invoiceNumber ?? null,
      invoice_date: normalized.invoiceDate ? new Date(toDateOnly(normalized.invoiceDate)) : null,
      amount: toMoney(normalized.amount),
      tax_amount: toMoney(normalized.taxAmount),
      total_amount: toMoney(normalized.totalAmount),
      seller_name: normalized.sellerName ?? null,
      buyer_name: normalized.buyerName ?? null,
      ocr_last_error: OCR_DUPLICATE_ERROR,
      ocr_completed_at: now,
      ocr_attempts: { increment: 1 },
      ocr_locked_at: null,
      ocr_lock_owner: null,
    },
  });
}

async function markInvoiceAsAvailable(prisma, invoiceId, normalized, ocrResultId) {
  const now = new Date();
  await prisma.invoicePool.update({
    where: { id: invoiceId },
    data: {
      status: "available",
      ocr_result_id: ocrResultId,
      invoice_code: normalized.invoiceCode ?? null,
      invoice_number: normalized.invoiceNumber ?? null,
      invoice_date: normalized.invoiceDate ? new Date(toDateOnly(normalized.invoiceDate)) : null,
      amount: toMoney(normalized.amount),
      tax_amount: toMoney(normalized.taxAmount),
      total_amount: toMoney(normalized.totalAmount),
      seller_name: normalized.sellerName ?? null,
      buyer_name: normalized.buyerName ?? null,
      ocr_last_error: null,
      ocr_completed_at: now,
      ocr_attempts: { increment: 1 },
      ocr_locked_at: null,
      ocr_lock_owner: null,
    },
  });
}

async function markInvoiceAsFailed(prisma, claimedInvoice, errorMessage, config) {
  const now = new Date();
  const nextAttempts = Number(claimedInvoice.ocr_attempts || 0) + 1;
  const isTerminal = nextAttempts >= config.ocrMaxAttempts;

  await prisma.invoicePool.update({
    where: { id: claimedInvoice.id },
    data: {
      status: isTerminal ? "ocr_failed" : "pending_ocr",
      ocr_attempts: nextAttempts,
      ocr_last_error: errorMessage,
      ocr_completed_at: isTerminal ? now : null,
      ocr_locked_at: null,
      ocr_lock_owner: null,
    },
  });
}

async function processClaimedInvoice(prisma, config, claimedInvoice, workerId, logger) {
  try {
    const absolutePath = path.resolve(config.storageDir, claimedInvoice.storage_path);
    const fileBuffer = await fs.readFile(absolutePath);
    const imageBase64 = fileBuffer.toString("base64");

    const { ocrResultId, normalized } = await getOrCreateOcrResult(
      prisma,
      config,
      claimedInvoice,
      imageBase64,
    );

    const duplicate = await findDuplicateInvoice(
      prisma,
      claimedInvoice.id,
      normalized.invoiceCode,
      normalized.invoiceNumber,
    );

    if (duplicate) {
      await markInvoiceAsDuplicate(prisma, claimedInvoice.id, normalized, ocrResultId);
      logger.warn(
        {
          invoiceId: claimedInvoice.id,
          duplicateInvoiceId: duplicate.id,
        },
        "OCR invoice marked as duplicate",
      );
      return;
    }

    await markInvoiceAsAvailable(prisma, claimedInvoice.id, normalized, ocrResultId);
  } catch (error) {
    const message = trimErrorMessage(error);
    await markInvoiceAsFailed(prisma, claimedInvoice, message, config);
    logger.error({ err: error, invoiceId: claimedInvoice.id, workerId }, "OCR processing failed");
  }
}

async function runOcrOnce(prisma, config, logger, workerId = `ocr:${process.pid}`) {
  const claimed = await claimPendingInvoices(prisma, config, workerId);
  if (!claimed.length) {
    return { claimed: 0 };
  }

  for (const item of claimed) {
    await processClaimedInvoice(prisma, config, item, workerId, logger);
  }

  return { claimed: claimed.length };
}

function startOcrJob(fastify) {
  if (!["tencent", "mock"].includes(fastify.config.ocrProvider)) {
    fastify.log.info({ provider: fastify.config.ocrProvider }, "OCR job disabled");
    return null;
  }

  let running = false;
  const workerId = `ocr-worker:${process.pid}`;

  const task = cron.schedule(fastify.config.ocrPollCron, async () => {
    if (running) {
      return;
    }

    running = true;
    try {
      await runOcrOnce(fastify.prisma, fastify.config, fastify.log, workerId);
    } catch (error) {
      fastify.log.error({ err: error }, "OCR job tick failed");
    } finally {
      running = false;
    }
  });

  fastify.addHook("onClose", async () => {
    task.stop();
    task.destroy();
  });

  fastify.log.info({ cron: fastify.config.ocrPollCron }, "OCR job scheduled");
  return task;
}

module.exports = {
  OCR_DUPLICATE_ERROR,
  claimPendingInvoices,
  runOcrOnce,
  startOcrJob,
};