reloop-v2/backend/src/problem-sets/parsing/parse-image-based-exam.ts

import * as fs from 'fs';
import { execFileSync } from 'child_process';
import { parseAnswerTable } from './parse-answer-table';
import { parseProblemPaper } from './parse-problem-paper';
import { extractText } from './extract-text';
import { ocrAnswerTable, renderPdfPageToPng } from './ocr-fallback';
import { kiceMathStrategy } from './strategies/math';
import { ImageBasedParseOptions, ParseResult, ParsedProblem } from './types';

const PAGE_PROBLEM_RE = /^(\d{1,2})\.\s+/gm;

export async function parseImageBasedExam(
  options: ImageBasedParseOptions,
): Promise<ParseResult> {
  const strategy = kiceMathStrategy;
  const warnings: string[] = [];
  const pageCount = getPdfPageCount(options.paperPdfPath);
  const textResult = parseProblemPaper(extractText(options.paperPdfPath, 'raw'), strategy);
  const textProblemByNumber = new Map(
    textResult.problems.map((problem) => [problem.number, problem]),
  );

  warnings.push(...textResult.warnings);

  let answers: ParseResult['answers'] = [];
  if (options.answerPdfPath) {
    const answerText = extractText(options.answerPdfPath, 'layout');
    const answerResult = parseAnswerTable(answerText, strategy);
    answers = answerResult.answers;
    warnings.push(...answerResult.warnings);

    const expectedAnswerCount = options.expectedProblemCount ?? strategy.maxProblemNumber;
    if (answers.length < expectedAnswerCount) {
      const ocrAnswers = await ocrAnswerTable(options.answerPdfPath);
      const merged = new Map(answers.map((answer) => [answer.number, answer.answerNumber]));

      for (const answer of ocrAnswers) {
        if (!merged.has(answer.number)) {
          merged.set(answer.number, answer.answerNumber);
        }
      }

      if (merged.size > answers.length) {
        warnings.push(
          `math answer OCR fallback merged ${merged.size - answers.length} missing answers`,
        );
      }

      answers = Array.from(merged.entries())
        .sort((left, right) => left[0] - right[0])
        .map(([number, answerNumber]) => ({ number, answerNumber }));
    }
  }

  fs.mkdirSync(options.renderedImageDir, { recursive: true });

  const pageAssignments = buildProblemPageAssignments(
    options.paperPdfPath,
    pageCount,
    options.expectedProblemCount ?? strategy.maxProblemNumber,
  );

  const problems: ParsedProblem[] = [];
  const problemCount = options.expectedProblemCount ?? inferProblemCount(textProblemByNumber, strategy);
  const renderedPages = new Map<number, { imagePath: string; imageUrl: string }>();

  for (let problemNumber = 1; problemNumber <= problemCount; problemNumber++) {
    const pageNumber = pageAssignments.get(problemNumber) ?? 1;
    const renderedPage =
      renderedPages.get(pageNumber) ??
      (await renderAndRegisterPage(
        options.paperPdfPath,
        pageNumber,
        options.renderedImageDir,
        options.renderedImageBaseUrl,
      ));

    renderedPages.set(pageNumber, renderedPage);

    const textProblem = textProblemByNumber.get(problemNumber);
    const needsReviewReasons = Array.from(
      new Set([...(textProblem?.needsReviewReasons ?? []), 'image-based']),
    );

    problems.push({
      number: problemNumber,
      bodyText: textProblem?.bodyText ?? '',
      choices: textProblem?.choices ?? emptyChoices(),
      passageStart: textProblem?.passageStart,
      passageEnd: textProblem?.passageEnd,
      imageUrl: renderedPage.imageUrl,
      pageImageUrl: renderedPage.imageUrl,
      needsReview: true,
      needsReviewReasons,
    });
  }

  if (
    typeof options.expectedProblemCount === 'number' &&
    problems.length !== options.expectedProblemCount
  ) {
    warnings.push(
      `expected ${options.expectedProblemCount} problems, parsed ${problems.length}`,
    );
  }

  return {
    problems,
    passages: textResult.passages,
    answers,
    warnings: Array.from(new Set(warnings)),
  };
}

function getPdfPageCount(pdfPath: string): number {
  const output = execFileSync('pdfinfo', [pdfPath], {
    encoding: 'utf-8',
    maxBuffer: 1024 * 1024,
  });
  const match = output.match(/^Pages:\s+(\d+)/m);
  if (!match) {
    throw new Error(`failed to read page count from pdfinfo output for ${pdfPath}`);
  }

  return Number(match[1]);
}

function buildProblemPageAssignments(
  pdfPath: string,
  pageCount: number,
  expectedProblemCount: number,
): Map<number, number> {
  const anchors: Array<{ number: number; pageNumber: number }> = [];
  let lastNumber = 0;

  for (let pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
    const pageText = extractPageText(pdfPath, pageNumber);
    const numbers = collectPageProblemNumbers(pageText, expectedProblemCount);

    for (const number of numbers) {
      if (number <= lastNumber) {
        continue;
      }

      anchors.push({ number, pageNumber });
      lastNumber = number;
    }
  }

  const assignments = new Map<number, number>();
  if (anchors.length === 0) {
    for (let problemNumber = 1; problemNumber <= expectedProblemCount; problemNumber++) {
      assignments.set(problemNumber, 1);
    }
    return assignments;
  }

  const firstAnchor = anchors[0];
  for (let problemNumber = 1; problemNumber < firstAnchor.number; problemNumber++) {
    assignments.set(problemNumber, firstAnchor.pageNumber);
  }

  for (let index = 0; index < anchors.length; index++) {
    const current = anchors[index];
    const next = anchors[index + 1];
    const endNumber = Math.min(next?.number ? next.number - 1 : expectedProblemCount, expectedProblemCount);

    for (let problemNumber = current.number; problemNumber <= endNumber; problemNumber++) {
      assignments.set(problemNumber, current.pageNumber);
    }
  }

  const fallbackPage = anchors[anchors.length - 1]?.pageNumber ?? 1;
  for (let problemNumber = 1; problemNumber <= expectedProblemCount; problemNumber++) {
    if (!assignments.has(problemNumber)) {
      assignments.set(problemNumber, fallbackPage);
    }
  }

  return assignments;
}

function extractPageText(pdfPath: string, pageNumber: number): string {
  return execFileSync(
    'pdftotext',
    ['-layout', '-f', String(pageNumber), '-l', String(pageNumber), '-enc', 'UTF-8', pdfPath, '-'],
    {
      encoding: 'utf-8',
      maxBuffer: 20 * 1024 * 1024,
    },
  );
}

function collectPageProblemNumbers(pageText: string, maxProblemNumber: number): number[] {
  const numbers: number[] = [];
  const seen = new Set<number>();
  PAGE_PROBLEM_RE.lastIndex = 0;

  let match: RegExpExecArray | null;
  while ((match = PAGE_PROBLEM_RE.exec(pageText)) !== null) {
    const number = Number(match[1]);
    if (number < 1 || number > maxProblemNumber || seen.has(number)) {
      continue;
    }

    seen.add(number);
    numbers.push(number);
  }

  return numbers.sort((left, right) => left - right);
}

async function renderAndRegisterPage(
  pdfPath: string,
  pageNumber: number,
  renderedImageDir: string,
  renderedImageBaseUrl: string,
): Promise<{ imagePath: string; imageUrl: string }> {
  const imagePath = await renderPdfPageToPng(
    pdfPath,
    pageNumber,
    `${renderedImageDir}/page.png`,
  );

  return {
    imagePath,
    imageUrl: `${stripTrailingSlash(renderedImageBaseUrl)}/page-${pageNumber}.png`,
  };
}

function inferProblemCount(
  textProblemByNumber: Map<number, ParsedProblem>,
  strategy: typeof kiceMathStrategy,
): number {
  const parsedNumbers = Array.from(textProblemByNumber.keys());
  if (parsedNumbers.length === 0) {
    return strategy.maxProblemNumber;
  }

  return Math.max(...parsedNumbers);
}

function stripTrailingSlash(value: string): string {
  return value.replace(/\/+$/, '');
}

function emptyChoices(): Record<'1' | '2' | '3' | '4' | '5', string> {
  return {
    '1': '',
    '2': '',
    '3': '',
    '4': '',
    '5': '',
  };
}