pdf-processing-server/app/routers/replacement.py

from fastapi import APIRouter
from pydantic import BaseModel
from tempfile import TemporaryDirectory
from glob import iglob
from pdf.processor import replace_texts
from core.config import settings
import base64
import os
import zipfile
import json

router = APIRouter()

class ReplacementRequest(BaseModel):
    pdf_zip: str # Base64 encoded zip file
    replacement_data: str # Base64 encoded JSON string
    preserve_original_fonts: bool = True

@router.post("/replacement")
async def replacement(request: ReplacementRequest):
    base_dir = settings.BASE_DIR
    pdf_zip = base64.b64decode(request.pdf_zip)
    replacement_data = json.loads(base64.b64decode(request.replacement_data).decode('utf-8'))
    with TemporaryDirectory(dir=base_dir, prefix=f'tmp_', ignore_cleanup_errors=True) as temp_dir:
        # Save the zip file to a temporary directory
        pdf_zip_path = os.path.join(temp_dir, 'pdf.zip')
        with open(pdf_zip_path, 'wb') as f:
            f.write(pdf_zip)

        # Extract the PDFs from the zip file
        pdf_dir = os.path.join(temp_dir, 'pdf')
        output_dir = os.path.join(temp_dir, 'output')
        os.makedirs(output_dir, exist_ok=True)
        with zipfile.ZipFile(pdf_zip_path, 'r') as zip_ref:
            zip_ref.extractall(pdf_dir)

        # Replace text in the PDFs
        pattern = os.path.join(pdf_dir, '*.pdf')
        for pdf_path in iglob(pattern):
            pdf_name = os.path.basename(pdf_path)
            output_path = os.path.join(output_dir, pdf_name)
            replace_texts(pdf_path, output_path, replacement_data[pdf_name], request.preserve_original_fonts)

        # Zip the output PDFs
        output_zip_path = os.path.join(temp_dir, 'output.zip')
        with zipfile.ZipFile(output_zip_path, 'w') as zip_ref:
            for output_pdf_path in iglob(os.path.join(output_dir, '*.pdf')):
                zip_ref.write(output_pdf_path, os.path.basename(output_pdf_path))
        zip_file = open(output_zip_path, 'rb')
        data = zip_file.read()
        zip_file.close()
        return {
            'status': 'success',
            'data': base64.b64encode(data).decode('utf-8'),
        }