app.py

import ctypes
import io
import uuid
from pathlib import Path

from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse, Response
from fastapi.staticfiles import StaticFiles
from PIL import Image

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

UPLOAD_DIR = Path("uploads")
UPLOAD_DIR.mkdir(exist_ok=True)

_lib = ctypes.CDLL("./librenderbench.so")

# int render_pdfium_page(path, page, scale, &buf, &w, &h, &stride, &ms)
_lib.render_pdfium_page.restype = ctypes.c_int
_lib.render_pdfium_page.argtypes = [
    ctypes.c_char_p,  # path
    ctypes.c_int,  # page_index
    ctypes.c_float,  # scale
    ctypes.c_int,  # grayscale
    ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)),  # out_buf
    ctypes.POINTER(ctypes.c_int),  # out_width
    ctypes.POINTER(ctypes.c_int),  # out_height
    ctypes.POINTER(ctypes.c_int),  # out_stride
    ctypes.POINTER(ctypes.c_double),  # out_render_ms
]

# int render_mupdf_page(path, page, scale, grayscale, &buf, &w, &h, &stride, &ms)
_lib.render_mupdf_page.restype = ctypes.c_int
_lib.render_mupdf_page.argtypes = [
    ctypes.c_char_p,
    ctypes.c_int,
    ctypes.c_float,
    ctypes.c_int,  # grayscale
    ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)),
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_double),
]

_lib.free_buffer.restype = None
_lib.free_buffer.argtypes = [ctypes.POINTER(ctypes.c_ubyte)]

# int get_page_info(path, page, &count, &width_pt, &height_pt)
_lib.get_page_info.restype = ctypes.c_int
_lib.get_page_info.argtypes = [
    ctypes.c_char_p,
    ctypes.c_int,
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_float),
    ctypes.POINTER(ctypes.c_float),
]

# In-memory store of uploaded PDFs: {file_id: filepath}
pdf_store: dict[str, Path] = {}


def _get_page_info(filepath: Path, page: int = 0):
    page_count = ctypes.c_int()
    width_pt = ctypes.c_float()
    height_pt = ctypes.c_float()
    rc = _lib.get_page_info(
        str(filepath).encode(),
        page,
        ctypes.byref(page_count),
        ctypes.byref(width_pt),
        ctypes.byref(height_pt),
    )
    if rc != 0:
        raise RuntimeError(f"get_page_info failed with code {rc}")
    return page_count.value, width_pt.value, height_pt.value


def _render_native(renderer: str, filepath: Path, page: int, ppi: float, grayscale: bool = False):
    scale = ppi / 72.0

    buf_ptr = ctypes.POINTER(ctypes.c_ubyte)()
    width = ctypes.c_int()
    height = ctypes.c_int()
    stride = ctypes.c_int()
    render_ms = ctypes.c_double()

    render_fn = (
        _lib.render_pdfium_page if renderer == "pdfium" else _lib.render_mupdf_page
    )

    rc = render_fn(
        str(filepath).encode(),
        page,
        ctypes.c_float(scale),
        ctypes.c_int(int(grayscale)),
        ctypes.byref(buf_ptr),
        ctypes.byref(width),
        ctypes.byref(height),
        ctypes.byref(stride),
        ctypes.byref(render_ms),
    )
    if rc != 0:
        raise RuntimeError(f"{renderer} render failed with code {rc}")

    w = width.value
    h = height.value
    s = stride.value

    # Copy pixel data out of the C-allocated buffer, then free it.
    buf_size = s * h
    pixel_data = bytes(
        ctypes.cast(buf_ptr, ctypes.POINTER(ctypes.c_ubyte * buf_size)).contents
    )
    _lib.free_buffer(buf_ptr)

    # Convert raw pixels → PNG (this is intentionally NOT timed).
    mode = "L" if grayscale else "RGB"
    img = Image.frombuffer(mode, (w, h), pixel_data, "raw", mode, s, 1)

    png_buf = io.BytesIO()
    img.save(png_buf, format="PNG")
    return png_buf.getvalue(), render_ms.value


@app.post("/upload")
async def upload_pdf(file: UploadFile = File(...)):
    file_id = uuid.uuid4().hex[:12]
    dest = UPLOAD_DIR / f"{file_id}.pdf"
    content = await file.read()
    dest.write_bytes(content)
    pdf_store[file_id] = dest

    page_count, _, _ = _get_page_info(dest)
    return {"file_id": file_id, "page_count": page_count, "filename": file.filename}


@app.get("/render/{file_id}")
def render_page(
    file_id: str,
    renderer: str = "pdfium",
    page: int = 0,
    ppi: int = 150,
    grayscale: bool = False,
):
    filepath = pdf_store.get(file_id)
    if not filepath or not filepath.exists():
        return Response(status_code=404, content="PDF not found")

    _, page_width_pt, page_height_pt = _get_page_info(filepath, page)

    img_bytes, elapsed = _render_native(renderer, filepath, page, ppi, grayscale)

    return Response(
        content=img_bytes,
        media_type="image/png",
        headers={
            "X-Render-Time-Ms": f"{elapsed:.2f}",
            "X-Page-Width-Pt": f"{page_width_pt:.1f}",
            "X-Page-Height-Pt": f"{page_height_pt:.1f}",
        },
    )


@app.get("/", response_class=HTMLResponse)
def index():
    return Path("static/index.html").read_text()