OCRController.java

package stirling.software.SPDF.controller.api.misc;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

import javax.imageio.ImageIO;

import org.apache.pdfbox.io.IOUtils;
import org.apache.pdfbox.multipdf.PDFMergerUtility;
import org.apache.pdfbox.pdfwriter.compress.CompressParameters;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.ModelAttribute;
import org.springframework.web.multipart.MultipartFile;

import io.github.pixee.security.Filenames;
import io.swagger.v3.oas.annotations.Operation;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import stirling.software.SPDF.config.EndpointConfiguration;
import stirling.software.SPDF.model.api.misc.ProcessPdfWithOcrRequest;
import stirling.software.common.annotations.AutoJobPostMapping;
import stirling.software.common.annotations.api.MiscApi;
import stirling.software.common.configuration.RuntimePathConfig;
import stirling.software.common.enumeration.ResourceWeight;
import stirling.software.common.model.ApplicationProperties;
import stirling.software.common.model.tool.ToolArity;
import stirling.software.common.model.tool.ToolFormat;
import stirling.software.common.model.tool.ToolIO;
import stirling.software.common.model.tool.ToolIOCase;
import stirling.software.common.model.tool.ToolIOWhen;
import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.util.ExceptionUtils;
import stirling.software.common.util.GeneralUtils;
import stirling.software.common.util.ProcessExecutor;
import stirling.software.common.util.ProcessExecutor.ProcessExecutorResult;
import stirling.software.common.util.TempDirectory;
import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.WebResponseUtils;

@MiscApi
@Slf4j
@RequiredArgsConstructor
public class OCRController {

    // Tesseract's recommended minimum; more pixels cost time without improving recognition
    private static final int OCR_RENDER_DPI = 300;

    private final ApplicationProperties applicationProperties;
    private final CustomPDFDocumentFactory pdfDocumentFactory;
    private final TempFileManager tempFileManager;
    private final EndpointConfiguration endpointConfiguration;
    private final RuntimePathConfig runtimePathConfig;

    private boolean isOcrMyPdfEnabled() {
        return endpointConfiguration.isGroupEnabled("OCRmyPDF");
    }

    private boolean isTesseractEnabled() {
        return endpointConfiguration.isGroupEnabled("tesseract");
    }

    /** Gets the list of available Tesseract languages from the tessdata directory */
    public List<String> getAvailableTesseractLanguages() {
        String tessdataDir = runtimePathConfig.getTessDataPath();
        File[] files = new File(tessdataDir).listFiles();
        if (files == null) {
            return Collections.emptyList();
        }
        return Arrays.stream(files)
                .filter(file -> file.getName().endsWith(".traineddata"))
                .map(file -> file.getName().replace(".traineddata", ""))
                .filter(lang -> !"osd".equalsIgnoreCase(lang))
                .toList();
    }

    @AutoJobPostMapping(
            consumes = MediaType.MULTIPART_FORM_DATA_VALUE,
            value = "/ocr-pdf",
            resourceWeight = ResourceWeight.LARGE_WEIGHT)
    @ToolIO(
            produces = ToolFormat.PDF,
            cases =
                    @ToolIOCase(
                            when = @ToolIOWhen(param = "sidecar", matches = "true"),
                            produces = ToolFormat.ZIP,
                            arity = ToolArity.SISO))
    @Operation(
            summary = "Process a PDF file with OCR",
            description =
                    "This endpoint processes a PDF file using OCR (Optical Character Recognition)."
                            + " Users can specify languages, sidecar, deskew, clean, cleanFinal, ocrType,"
                            + " ocrRenderType, and removeImagesAfter options. Uses OCRmyPDF if available,"
                            + " falls back to Tesseract.")
    public ResponseEntity<Resource> processPdfWithOCR(
            @ModelAttribute ProcessPdfWithOcrRequest request)
            throws IOException, InterruptedException {
        MultipartFile inputFile = request.getFileInput();
        List<String> selectedLanguages = request.getLanguages();
        boolean sidecar = request.isSidecar();
        Boolean deskew = request.isDeskew();
        Boolean rotatePages = request.isRotatePages();
        Boolean clean = request.isClean();
        Boolean cleanFinal = request.isCleanFinal();
        String ocrType = request.getOcrType();
        String ocrRenderType = request.getOcrRenderType();
        Boolean removeImagesAfter = request.isRemoveImagesAfter();

        if (selectedLanguages == null || selectedLanguages.isEmpty()) {
            throw ExceptionUtils.createOcrLanguageRequiredException();
        }

        if (!"hocr".equals(ocrRenderType) && !"sandwich".equals(ocrRenderType)) {
            throw ExceptionUtils.createOcrInvalidRenderTypeException();
        }

        // Get available Tesseract languages
        List<String> availableLanguages = getAvailableTesseractLanguages();

        // Validate selected languages
        selectedLanguages =
                selectedLanguages.stream().filter(availableLanguages::contains).toList();

        if (selectedLanguages.isEmpty()) {
            throw ExceptionUtils.createOcrInvalidLanguagesException();
        }

        TempFile tempOutputFile = new TempFile(tempFileManager, ".pdf");
        TempFile tempZipFile = null;
        boolean pdfOwnershipTransferred = false;
        boolean zipOwnershipTransferred = false;
        try (TempFile tempInputFile = new TempFile(tempFileManager, ".pdf");
                TempFile sidecarTextFile = sidecar ? new TempFile(tempFileManager, ".txt") : null) {

            inputFile.transferTo(tempInputFile.getFile());

            // Use OCRmyPDF if available (no fallback - error if it fails)
            if (isOcrMyPdfEnabled()) {
                processWithOcrMyPdf(
                        selectedLanguages,
                        sidecar,
                        deskew,
                        rotatePages,
                        clean,
                        cleanFinal,
                        ocrType,
                        ocrRenderType,
                        removeImagesAfter,
                        tempInputFile.getPath(),
                        tempOutputFile.getPath(),
                        sidecarTextFile != null ? sidecarTextFile.getPath() : null);
                log.info("OCRmyPDF processing completed successfully");
            }
            // Use Tesseract only if OCRmyPDF is not available
            else if (isTesseractEnabled()) {
                processWithTesseract(
                        selectedLanguages,
                        ocrType,
                        tempInputFile.getPath(),
                        tempOutputFile.getPath());
                log.info("Tesseract processing completed successfully");
            } else {
                throw ExceptionUtils.createOcrToolsUnavailableException();
            }

            // Return the OCR processed PDF as a response
            String outputFilename =
                    GeneralUtils.removeExtension(
                                    Filenames.toSimpleFileName(inputFile.getOriginalFilename()))
                            + "_OCR.pdf";

            if (sidecar && sidecarTextFile != null) {
                // Create a zip file containing both the PDF and the text file
                String outputZipFilename =
                        GeneralUtils.removeExtension(
                                        Filenames.toSimpleFileName(inputFile.getOriginalFilename()))
                                + "_OCR.zip";

                tempZipFile = new TempFile(tempFileManager, ".zip");
                try (ZipOutputStream zipOut =
                        new ZipOutputStream(Files.newOutputStream(tempZipFile.getPath()))) {

                    // Add PDF file to the zip
                    ZipEntry pdfEntry = new ZipEntry(outputFilename);
                    zipOut.putNextEntry(pdfEntry);
                    Files.copy(tempOutputFile.getPath(), zipOut);
                    zipOut.closeEntry();

                    // Add text file to the zip
                    ZipEntry txtEntry = new ZipEntry(outputFilename.replace(".pdf", ".txt"));
                    zipOut.putNextEntry(txtEntry);
                    Files.copy(sidecarTextFile.getPath(), zipOut);
                    zipOut.closeEntry();

                    zipOut.finish();
                }

                // The intermediate PDF temp file is no longer needed; only the zip is streamed.
                tempOutputFile.close();
                pdfOwnershipTransferred = true;
                ResponseEntity<Resource> response =
                        WebResponseUtils.fileToWebResponse(
                                tempZipFile, outputZipFilename, MediaType.APPLICATION_OCTET_STREAM);
                zipOwnershipTransferred = true;
                return response;
            } else {
                ResponseEntity<Resource> response =
                        WebResponseUtils.pdfFileToWebResponse(tempOutputFile, outputFilename);
                pdfOwnershipTransferred = true;
                return response;
            }
        } finally {
            if (!pdfOwnershipTransferred) {
                tempOutputFile.close();
            }
            if (tempZipFile != null && !zipOwnershipTransferred) {
                tempZipFile.close();
            }
        }
    }

    private void processWithOcrMyPdf(
            List<String> selectedLanguages,
            Boolean sidecar,
            Boolean deskew,
            Boolean rotatePages,
            Boolean clean,
            Boolean cleanFinal,
            String ocrType,
            String ocrRenderType,
            Boolean removeImagesAfter,
            Path tempInputFile,
            Path tempOutputFile,
            Path sidecarTextPath)
            throws IOException, InterruptedException {

        // Build OCRmyPDF command
        String languageOption = String.join("+", selectedLanguages);

        List<String> command =
                new ArrayList<>(
                        Arrays.asList(
                                runtimePathConfig.getOcrMyPdfPath(),
                                "--verbose",
                                "2",
                                "--output-type",
                                "pdf",
                                "--pdf-renderer",
                                ocrRenderType));

        if (sidecar != null && sidecar && sidecarTextPath != null) {
            command.add("--sidecar");
            command.add(sidecarTextPath.toString());
        }

        if (deskew != null && deskew) {
            command.add("--deskew");
        }
        if (rotatePages != null && rotatePages) {
            // Tesseract OSD-based automatic page orientation correction (90/180/270)
            command.add("--rotate-pages");
        }
        if (clean != null && clean) {
            command.add("--clean");
        }
        if (cleanFinal != null && cleanFinal) {
            command.add("--clean-final");
        }
        if (ocrType != null && !ocrType.isEmpty()) {
            if ("force-ocr".equals(ocrType)) {
                command.add("--force-ocr");
            } else {
                // Default for 'Normal' and 'skip-text': use --skip-text
                // ocrmypdf 17+ requires explicit flag when PDF already contains text
                command.add("--skip-text");
            }
        }
        command.add("--invalidate-digital-signatures");

        command.addAll(
                Arrays.asList(
                        "--language",
                        languageOption,
                        tempInputFile.toString(),
                        tempOutputFile.toString()));

        // Run CLI command
        ProcessExecutorResult result =
                ProcessExecutor.getInstance(ProcessExecutor.Processes.OCR_MY_PDF)
                        .runCommandWithOutputHandling(command);

        if (result.getRc() != 0
                && result.getMessages().contains("multiprocessing/synchronize.py")
                && result.getMessages().contains("OSError: [Errno 38] Function not implemented")) {
            command.add("--jobs");
            command.add("1");
            result =
                    ProcessExecutor.getInstance(ProcessExecutor.Processes.OCR_MY_PDF)
                            .runCommandWithOutputHandling(command);
        }

        if (result.getRc() != 0) {
            throw ExceptionUtils.createOcrProcessingFailedException(result.getRc());
        }

        // Remove images from the OCR processed PDF if the flag is set to true
        if (removeImagesAfter != null && removeImagesAfter) {
            try (TempFile tempPdfWithoutImages = new TempFile(tempFileManager, "_no_images.pdf")) {
                List<String> gsCommand =
                        Arrays.asList(
                                "gs",
                                "-sDEVICE=pdfwrite",
                                "-dFILTERIMAGE",
                                "-o",
                                tempPdfWithoutImages.getPath().toString(),
                                tempOutputFile.toString());

                ProcessExecutor.getInstance(ProcessExecutor.Processes.GHOSTSCRIPT)
                        .runCommandWithOutputHandling(gsCommand);

                // Replace output file with version without images
                Files.copy(
                        tempPdfWithoutImages.getPath(),
                        tempOutputFile,
                        java.nio.file.StandardCopyOption.REPLACE_EXISTING);
            }
        }
    }

    private void processWithTesseract(
            List<String> selectedLanguages, String ocrType, Path tempInputFile, Path tempOutputFile)
            throws IOException, InterruptedException {

        // Create temp directory for Tesseract processing
        try (TempDirectory tempDir = new TempDirectory(tempFileManager)) {
            File tempOutputDir = new File(tempDir.getPath().toFile(), "output");
            File tempImagesDir = new File(tempDir.getPath().toFile(), "images");
            File finalOutputFile = new File(tempDir.getPath().toFile(), "final_output.pdf");

            // Create directories
            tempOutputDir.mkdirs();
            tempImagesDir.mkdirs();

            PDFMergerUtility merger = new PDFMergerUtility();
            merger.setDestinationFileName(finalOutputFile.toString());

            try (PDDocument document = pdfDocumentFactory.load(tempInputFile.toFile())) {
                PDFRenderer pdfRenderer = new PDFRenderer(document);
                pdfRenderer.setSubsamplingAllowed(
                        true); // Enable subsampling to reduce memory usage
                int pageCount = document.getNumberOfPages();

                for (int pageNum = 0; pageNum < pageCount; pageNum++) {
                    PDPage page = document.getPage(pageNum);
                    boolean hasText;

                    // Check for existing text
                    try (PDDocument tempDoc = new PDDocument()) {
                        tempDoc.addPage(page);
                        PDFTextStripper stripper = new PDFTextStripper();
                        hasText = !stripper.getText(tempDoc).trim().isEmpty();
                    }

                    boolean shouldOcr =
                            switch (ocrType) {
                                case "skip-text" -> !hasText;
                                case "force-ocr" -> true;
                                default -> true;
                            };

                    File pageOutputPath =
                            new File(
                                    tempOutputDir,
                                    String.format(Locale.ROOT, "page_%d.pdf", pageNum));

                    if (shouldOcr) {
                        // Convert page to image
                        BufferedImage image;

                        // maxDPI is a safety ceiling for user-supplied values, not a target
                        int renderDpi = OCR_RENDER_DPI;
                        if (applicationProperties != null
                                && applicationProperties.getSystem() != null) {
                            renderDpi =
                                    Math.min(
                                            renderDpi,
                                            applicationProperties.getSystem().getMaxDPI());
                        }
                        final int dpi = renderDpi;
                        final int currentPageNum = pageNum;

                        image =
                                ExceptionUtils.handleOomRendering(
                                        currentPageNum + 1,
                                        dpi,
                                        () -> pdfRenderer.renderImageWithDPI(currentPageNum, dpi));
                        File imagePath =
                                new File(
                                        tempImagesDir,
                                        String.format(Locale.ROOT, "page_%d.png", pageNum));
                        ImageIO.write(image, "png", imagePath);

                        // Build OCR command
                        List<String> command = new ArrayList<>();
                        command.add("tesseract");
                        command.add(imagePath.toString());
                        String outputBase =
                                new File(
                                                tempOutputDir,
                                                String.format(Locale.ROOT, "page_%d", pageNum))
                                        .toString();
                        command.add(outputBase);
                        command.add("-l");
                        command.add(String.join("+", selectedLanguages));
                        command.add("pdf"); // Always output PDF

                        ProcessExecutorResult result =
                                ProcessExecutor.getInstance(ProcessExecutor.Processes.TESSERACT)
                                        .runCommandWithOutputHandling(command);

                        if (result.getRc() != 0) {
                            throw ExceptionUtils.createRuntimeException(
                                    "error.commandFailed",
                                    "{0} command failed with exit code: {1}",
                                    null,
                                    "Tesseract",
                                    result.getRc());
                        }

                        // Verify the OCR'd PDF was created
                        if (!pageOutputPath.exists()) {
                            log.warn(
                                    "Tesseract did not create expected output file: {}. Page may be blank or unreadable.",
                                    pageOutputPath.getAbsolutePath());
                            // Save original page without OCR as fallback
                            try (PDDocument pageDoc = new PDDocument()) {
                                pageDoc.addPage(page);
                                // NO_COMPRESSION: page is copied from another document;
                                // PDFBox 3.0.7 compressed writer (PDFBOX-6203) drops shared
                                // resources, corrupting fonts. Revert once on 3.0.8.
                                pageDoc.save(pageOutputPath, CompressParameters.NO_COMPRESSION);
                            }
                        }

                        // Add OCR'd PDF to merger
                        merger.addSource(pageOutputPath);
                    } else {
                        // Save original page without OCR
                        try (PDDocument pageDoc = new PDDocument()) {
                            pageDoc.addPage(page);
                            // NO_COMPRESSION: page is copied from another document; PDFBox 3.0.7
                            // compressed writer (PDFBOX-6203) drops shared resources, corrupting
                            // fonts on retained text pages. Revert once on 3.0.8.
                            pageDoc.save(pageOutputPath, CompressParameters.NO_COMPRESSION);
                            merger.addSource(pageOutputPath);
                        }
                    }
                }
            }

            // Merge all pages into final PDF
            merger.mergeDocuments(IOUtils.createTempFileOnlyStreamCache());

            // Copy final output to the expected location
            Files.copy(
                    finalOutputFile.toPath(),
                    tempOutputFile,
                    java.nio.file.StandardCopyOption.REPLACE_EXISTING);
        }
    }
}