Voiced/src/pdf/flatten.ts

/** The slice of pdf.js the flatten needs, so Node tests and the webview can supply their own. */
export interface PageViewportLike {
  width: number;
  height: number;
}

export interface RenderTaskLike {
  promise: Promise<unknown>;
  cancel?: () => void;
}

export interface PageLike {
  getViewport(params: { scale: number }): PageViewportLike;
  render(params: { canvasContext: unknown; canvas: unknown; viewport: PageViewportLike; intent?: string }): RenderTaskLike;
  cleanup?: () => unknown;
}

export interface PdfDocLike {
  numPages: number;
  getPage(n: number): Promise<PageLike>;
}

export interface CanvasLike {
  width: number;
  height: number;
  getContext(kind: "2d"): unknown;
}

export interface FlatPage {
  jpeg: Uint8Array;
  widthPt: number;
  heightPt: number;
  pxWidth: number;
  pxHeight: number;
  renderMs: number;
  encodeMs: number;
}

export interface FlattenResult {
  pages: FlatPage[];
}

export interface FlattenOptions {
  pdf: PdfDocLike;
  dpi?: number;
  quality?: number;
  createCanvas: (width: number, height: number) => CanvasLike;
  encodeJpeg: (canvas: CanvasLike, quality: number) => Promise<Uint8Array>;
  signal?: AbortSignal;
  /** Called as each page starts: (page being rendered, page count). */
  onProgress?: (page: number, pages: number) => void;
  /** Awaited after each page so the UI can paint and handle input between the long tasks. */
  yieldToUi?: () => Promise<void>;
  now?: () => number;
}

export class FlattenCancelled extends Error {
  constructor() {
    super("The flatten was cancelled");
    this.name = "FlattenCancelled";
  }
}

export class FlattenError extends Error {
  readonly code = "FLATTEN_FAILED";
  constructor(
    readonly page: number,
    readonly reason: string,
  ) {
    super(page > 0 ? `Could not flatten page ${page}: ${reason}` : `Could not flatten the PDF: ${reason}`);
    this.name = "FlattenError";
  }
}

const reasonOf = (e: unknown) => (e instanceof Error ? e.message : String(e));

/**
 * Rasterises every page to a JPEG at `dpi` (default 300) on a white background. One canvas lives at a
 * time and is released before the next page. There is deliberately no fallback to a lower DPI: a canvas
 * or encode failure throws a FlattenError and the caller decides what to offer.
 */
export async function flattenPages(opts: FlattenOptions): Promise<FlattenResult> {
  const { pdf, dpi = 300, quality = 0.92, createCanvas, encodeJpeg, signal, onProgress, yieldToUi, now = () => performance.now() } = opts;
  const scale = dpi / 72;
  const total = pdf.numPages;
  if (!(total > 0)) throw new FlattenError(0, "the PDF has no pages");

  const checkAbort = () => {
    if (signal?.aborted) throw new FlattenCancelled();
  };

  const pages: FlatPage[] = [];
  for (let n = 1; n <= total; n++) {
    checkAbort();
    onProgress?.(n, total);
    checkAbort(); // onProgress may be where the caller cancels

    let canvas: CanvasLike | null = null;
    let page: PageLike | null = null;
    try {
      page = await pdf.getPage(n);
      const unit = page.getViewport({ scale: 1 });
      const viewport = page.getViewport({ scale });
      // The epsilon keeps 792 pt x 300/72 (3300.0000000000005) from becoming a 3301 px page.
      const pxWidth = Math.ceil(viewport.width - 1e-6);
      const pxHeight = Math.ceil(viewport.height - 1e-6);

      canvas = createCanvas(pxWidth, pxHeight);
      const ctx = canvas.getContext("2d") as { fillStyle: unknown; fillRect(x: number, y: number, w: number, h: number): void } | null;
      if (!ctx) throw new Error("the canvas has no 2D context");
      ctx.fillStyle = "#ffffff";
      ctx.fillRect(0, 0, pxWidth, pxHeight);

      checkAbort();
      const t0 = now();
      const task = page.render({ canvasContext: ctx, canvas, viewport, intent: "print" });
      const onAbort = () => task.cancel?.();
      signal?.addEventListener("abort", onAbort, { once: true });
      try {
        await task.promise;
      } finally {
        signal?.removeEventListener("abort", onAbort);
      }
      checkAbort();
      const t1 = now();
      const jpeg = await encodeJpeg(canvas, quality);
      const t2 = now();
      if (!jpeg || jpeg.byteLength < 4) throw new Error("the JPEG encoder returned no data");
      pages.push({
        jpeg,
        widthPt: unit.width,
        heightPt: unit.height,
        pxWidth,
        pxHeight,
        renderMs: Math.round(t1 - t0),
        encodeMs: Math.round(t2 - t1),
      });
    } catch (e) {
      if (e instanceof FlattenCancelled || signal?.aborted) throw new FlattenCancelled();
      if (e instanceof FlattenError) throw e;
      throw new FlattenError(n, reasonOf(e));
    } finally {
      if (canvas) {
        // Releases the backing store now rather than at garbage collection (2480 x 3508 x 4 = 35 MB).
        canvas.width = 0;
        canvas.height = 0;
      }
      try {
        page?.cleanup?.();
      } catch {
        /* cleanup is best effort */
      }
    }
    if (yieldToUi) await yieldToUi();
  }
  return { pages };
}