reloop-v2/backend/scripts/ocr-fix.ts

import * as fs from 'fs';
import * as path from 'path';
import { execFileSync } from 'child_process';
import { PrismaClient, Prisma } from '@prisma/client';
import { ocrAnswerTable, ocrProblemPage } from '../src/problem-sets/parsing/ocr-fallback';

type TargetSubject = '수학';
type ChoiceKey = '1' | '2' | '3' | '4' | '5';

interface CliOptions {
  dryRun: boolean;
  year?: number;
  subject?: string;
  answersOnly: boolean;
}

interface SubjectFiles {
  subject: TargetSubject;
  answerPdfPath: string;
  problemPdfPath: string;
}

interface ProblemSnapshot {
  id: number;
  number: number;
  bodyText: string | null;
  choices: Prisma.JsonValue | null;
  answerNumber: number | null;
  needsReview: boolean;
}

interface SetSummary {
  year: number;
  subject: TargetSubject;
  fixedAnswers: number;
  fixedProblems: number;
  skipped: number;
}

interface AnswerUpdate {
  problemId: number;
  number: number;
  answerNumber: number;
  needsReview: boolean;
  shouldWrite: boolean;
}

interface ProblemUpdate {
  kind: 'update';
  problemId: number;
  number: number;
  bodyText: string;
  choices: Record<ChoiceKey, string>;
  needsReview: boolean;
  shouldWrite: boolean;
}

interface ProblemCreate {
  kind: 'create';
  number: number;
  title: string;
  difficulty: number;
  bodyText: string;
  choices: Record<ChoiceKey, string>;
  answerNumber: number | null;
  needsReview: boolean;
  shouldWrite: boolean;
}

const prisma = new PrismaClient();
const DATA_DIR = path.resolve(__dirname, '../../data/kice');
const TARGET_SUBJECTS: TargetSubject[] = ['수학'];
const EXPECTED_PROBLEM_COUNT: Record<TargetSubject, number> = {
  수학: 30,
};

async function main(): Promise<void> {
  const options = parseArgs();
  const years = options.year ? [options.year] : [2025];
  const subjects = options.subject ? [options.subject] : TARGET_SUBJECTS;
  const summaries: SetSummary[] = [];

  try {
    for (const year of years) {
      for (const subjectName of subjects) {
        if (!TARGET_SUBJECTS.includes(subjectName as TargetSubject)) {
          console.warn(`skip unsupported subject: ${subjectName}`);
          continue;
        }

        const subject = subjectName as TargetSubject;
        const files = resolveSubjectFiles(year, subject);
        const problemSet = await prisma.problemSet.findFirst({
          where: { year, examType: 'sat', subjectName: subject },
          include: { problems: { orderBy: { number: 'asc' } } },
        });

        if (!problemSet) {
          console.warn(`skip ${year} ${subject}: problem set not found in DB`);
          continue;
        }

        const summary = await processProblemSet(problemSet, files, options);
        summaries.push(summary);
        console.log(
          `${year} ${subject}: fixed ${summary.fixedAnswers} answers, fixed ${summary.fixedProblems} problems, skipped ${summary.skipped}`,
        );
      }
    }
  } finally {
    await prisma.$disconnect();
  }

  if (summaries.length === 0) {
    console.log('no matching problem sets processed');
  }
}

function parseArgs(): CliOptions {
  const args = process.argv.slice(2);
  const options: CliOptions = {
    dryRun: false,
    answersOnly: false,
  };

  for (const arg of args) {
    if (arg === '--dry-run') {
      options.dryRun = true;
      continue;
    }
    if (arg === '--answers-only') {
      options.answersOnly = true;
      continue;
    }
    if (arg.startsWith('--year=')) {
      options.year = Number(arg.split('=')[1]);
      continue;
    }
    if (arg.startsWith('--subject=')) {
      options.subject = arg.split('=')[1];
    }
  }

  return options;
}

async function processProblemSet(
  problemSet: {
    id: number;
    year: number;
    subjectName: string;
    problems: ProblemSnapshot[];
  },
  files: SubjectFiles,
  options: CliOptions,
): Promise<SetSummary> {
  const answerUpdates = await collectAnswerUpdates(problemSet.problems, files.answerPdfPath);
  const problemUpdates: Array<ProblemUpdate | ProblemCreate> = options.answersOnly
    ? []
    : await collectProblemUpdates(
        problemSet.year,
        files.subject,
        files.problemPdfPath,
        problemSet.problems,
      );

  const skipped =
    answerUpdates.filter((update) => !update.shouldWrite).length +
    problemUpdates.filter((update) => !update.shouldWrite).length;

  if (options.dryRun) {
    printDryRun(problemSet.year, files.subject, answerUpdates, problemUpdates);
    return {
      year: problemSet.year,
      subject: files.subject,
      fixedAnswers: answerUpdates.filter((update) => update.shouldWrite).length,
      fixedProblems: problemUpdates.filter((update) => update.shouldWrite).length,
      skipped,
    };
  }

  await prisma.$transaction(async (tx) => {
    for (const update of answerUpdates) {
      if (!update.shouldWrite) continue;
      await tx.problem.update({
        where: { id: update.problemId },
        data: {
          answerNumber: update.answerNumber,
          needsReview: update.needsReview,
        },
      });
    }

    for (const update of problemUpdates) {
      if (!update.shouldWrite) continue;

      if (update.kind === 'update') {
        await tx.problem.update({
          where: { id: update.problemId },
          data: {
            bodyText: update.bodyText,
            choices: update.choices as Prisma.InputJsonValue,
            needsReview: update.needsReview,
          },
        });
        continue;
      }

      await tx.problem.upsert({
        where: {
          problemSetId_number: {
            problemSetId: problemSet.id,
            number: update.number,
          },
        },
        update: {
          bodyText: update.bodyText,
          choices: update.choices as Prisma.InputJsonValue,
          answerNumber: update.answerNumber,
          needsReview: update.needsReview,
          title: update.title,
          difficulty: update.difficulty,
        },
        create: {
          problemSetId: problemSet.id,
          number: update.number,
          title: update.title,
          difficulty: update.difficulty,
          bodyText: update.bodyText,
          choices: update.choices as Prisma.InputJsonValue,
          answerNumber: update.answerNumber,
          needsReview: update.needsReview,
        },
      });
    }
  });

  return {
    year: problemSet.year,
    subject: files.subject,
    fixedAnswers: answerUpdates.filter((update) => update.shouldWrite).length,
    fixedProblems: problemUpdates.filter((update) => update.shouldWrite).length,
    skipped,
  };
}

async function collectAnswerUpdates(
  problems: ProblemSnapshot[],
  answerPdfPath: string,
): Promise<AnswerUpdate[]> {
  const allAnswersMissing = problems.length > 0 && problems.every((problem) => problem.answerNumber === null);
  if (!allAnswersMissing) {
    return [];
  }

  const answers = await ocrAnswerTable(answerPdfPath);
  const answersByNumber = new Map(answers.map((answer) => [answer.number, answer.answerNumber]));

  return problems.flatMap((problem) => {
    const answerNumber = answersByNumber.get(problem.number);
    if (!answerNumber) {
      return [];
    }

    const currentChoices = normalizeStoredChoices(problem.choices);
    const currentBodyText = (problem.bodyText || '').trim();
    const needsReview = computeNeedsReview(currentBodyText, currentChoices, answerNumber);

    return [
      {
        problemId: problem.id,
        number: problem.number,
        answerNumber,
        needsReview,
        shouldWrite: problem.answerNumber !== answerNumber || problem.needsReview !== needsReview,
      },
    ];
  });
}

async function collectProblemUpdates(
  year: number,
  subject: TargetSubject,
  problemPdfPath: string,
  problems: ProblemSnapshot[],
): Promise<ProblemUpdate[]> {
  const pageCount = getPdfPageCount(problemPdfPath);
  const reviewProblems = problems.filter((problem) => problem.needsReview);
  const updates: ProblemUpdate[] = [];

  for (const problem of reviewProblems) {
    const pageNumber =
      findProblemPageByText(problemPdfPath, pageCount, problem.number) ??
      estimateProblemPage(subject, problem.number, pageCount);
    const ocrProblems = await ocrProblemPage(problemPdfPath, pageNumber, [problem.number]);
    const ocrProblem = ocrProblems.find((entry) => entry.number === problem.number);

    if (!ocrProblem) {
      updates.push({
        kind: 'update',
        problemId: problem.id,
        number: problem.number,
        bodyText: (problem.bodyText || '').trim(),
        choices: normalizeStoredChoices(problem.choices),
        needsReview: problem.needsReview,
        shouldWrite: false,
      });
      continue;
    }

    const currentBodyText = (problem.bodyText || '').trim();
    const currentChoices = normalizeStoredChoices(problem.choices);
    const ocrChoices = normalizeChoiceRecord(ocrProblem.choices);
    const shouldWrite = isClearlyBetter(currentBodyText, currentChoices, ocrProblem.bodyText, ocrChoices);
    const answerNumber = coerceAnswerNumber(problem.answerNumber);
    const needsReview = computeNeedsReview(ocrProblem.bodyText, ocrChoices, answerNumber);

    updates.push({
      kind: 'update',
      problemId: problem.id,
      number: problem.number,
      bodyText: shouldWrite ? ocrProblem.bodyText : currentBodyText,
      choices: shouldWrite ? ocrChoices : currentChoices,
      needsReview: shouldWrite ? needsReview : problem.needsReview,
      shouldWrite,
    });
  }

  return updates;
}

function printDryRun(
  year: number,
  subject: TargetSubject,
  answerUpdates: AnswerUpdate[],
  problemUpdates: Array<ProblemUpdate | ProblemCreate>,
): void {
  for (const update of answerUpdates) {
    if (!update.shouldWrite) continue;
    console.log(
      `[dry-run] ${year} ${subject} answer #${update.number} -> ${update.answerNumber}, needsReview=${update.needsReview}`,
    );
  }

  for (const update of problemUpdates) {
    if (!update.shouldWrite) continue;
    console.log(
      `[dry-run] ${year} ${subject} problem #${update.number} OCR body=${truncate(update.bodyText, 80)}`,
    );
  }
}

function resolveSubjectFiles(year: number, subject: TargetSubject): SubjectFiles {
  const subjectDir = path.join(DATA_DIR, String(year), subject);
  const answerPdfPath = findFile(subjectDir, /_정답표?\.pdf$/);
  const problemPdfPath =
    findFile(subjectDir, /_문제지_홀수형\.pdf$/) ||
    findFile(subjectDir, /_문제지(_홀수형)?\.pdf$/) ||
    findFile(subjectDir, /_문제\.pdf$/);

  if (!answerPdfPath || !problemPdfPath) {
    throw new Error(`missing PDFs for ${year} ${subject}`);
  }

  return {
    subject,
    answerPdfPath,
    problemPdfPath,
  };
}

function findFile(dir: string, pattern: RegExp): string | null {
  if (!fs.existsSync(dir)) {
    return null;
  }

  const entry = fs.readdirSync(dir).find((value) => pattern.test(value));
  return entry ? path.join(dir, entry) : null;
}

function requireExistingFile(filePath: string): string {
  if (!fs.existsSync(filePath)) {
    throw new Error(`missing file: ${filePath}`);
  }
  return filePath;
}

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(`unable to determine page count for ${pdfPath}`);
  }
  return Number(match[1]);
}

function findProblemPageByText(
  pdfPath: string,
  pageCount: number,
  problemNumber: number,
): number | null {
  const pattern = new RegExp(`(^|\\s)${problemNumber}\\.\\s`, 'm');

  for (let pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
    const text = execFileSync(
      'pdftotext',
      ['-raw', '-enc', 'UTF-8', '-f', String(pageNumber), '-l', String(pageNumber), pdfPath, '-'],
      {
        encoding: 'utf-8',
        maxBuffer: 10 * 1024 * 1024,
      },
    );

    if (pattern.test(text)) {
      return pageNumber;
    }
  }

  return null;
}

function estimateProblemPage(
  _subject: TargetSubject,
  problemNumber: number,
  pageCount: number,
): number {
  const problemsPerPage = 4;
  return Math.max(1, Math.min(pageCount, Math.ceil(problemNumber / problemsPerPage)));
}

function normalizeStoredChoices(value: Prisma.JsonValue | null): Record<ChoiceKey, string> {
  if (!value || typeof value !== 'object' || Array.isArray(value)) {
    return emptyChoices();
  }

  const record = value as Record<string, unknown>;
  return {
    '1': typeof record['1'] === 'string' ? record['1'].trim() : '',
    '2': typeof record['2'] === 'string' ? record['2'].trim() : '',
    '3': typeof record['3'] === 'string' ? record['3'].trim() : '',
    '4': typeof record['4'] === 'string' ? record['4'].trim() : '',
    '5': typeof record['5'] === 'string' ? record['5'].trim() : '',
  };
}

function normalizeChoiceRecord(value: Record<ChoiceKey, string>): Record<ChoiceKey, string> {
  return {
    '1': value['1'].trim(),
    '2': value['2'].trim(),
    '3': value['3'].trim(),
    '4': value['4'].trim(),
    '5': value['5'].trim(),
  };
}

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

function countFilledChoices(choices: Record<ChoiceKey, string>): number {
  return (Object.keys(choices) as ChoiceKey[]).filter((key) => choices[key].length > 0).length;
}

function scoreProblem(bodyText: string, choices: Record<ChoiceKey, string>): number {
  const filledChoices = countFilledChoices(choices);
  const choiceChars = Object.values(choices).reduce((sum, value) => sum + value.length, 0);
  return filledChoices * 1_000 + choiceChars * 10 + bodyText.length;
}

function isClearlyBetter(
  currentBodyText: string,
  currentChoices: Record<ChoiceKey, string>,
  nextBodyText: string,
  nextChoices: Record<ChoiceKey, string>,
): boolean {
  const currentScore = scoreProblem(currentBodyText, currentChoices);
  const nextScore = scoreProblem(nextBodyText, nextChoices);
  return nextScore > currentScore + 100;
}

function computeNeedsReview(
  bodyText: string,
  choices: Record<ChoiceKey, string>,
  answerNumber: number | null,
): boolean {
  if (!answerNumber) {
    return true;
  }
  if (bodyText.trim().length < 10) {
    return true;
  }

  for (const key of ['1', '2', '3', '4', '5'] as const) {
    const value = choices[key].trim();
    if (value.length < 3) {
      return true;
    }
  }

  return false;
}

function coerceAnswerNumber(value: number | null): number | null {
  return typeof value === 'number' && Number.isInteger(value) ? value : null;
}

function truncate(value: string, length: number): string {
  return value.length <= length ? value : `${value.slice(0, length - 3)}...`;
}

void main().catch((error) => {
  console.error(error);
  process.exitCode = 1;
});