pdf-tools/src/pdf/pdfThumbnailService.ts

import * as pdfjsLib from 'pdfjs-dist';
import pdfjsWorker from 'pdfjs-dist/build/pdf.worker?worker&url';

// pdf.js worker setup for Vite
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(pdfjsLib as any).GlobalWorkerOptions.workerSrc = pdfjsWorker;

type RotationsMap = Record<number, number>; // key: 0-based page index, value: degrees

/**
 * Helper: create a *copy* of the data for pdf.js to consume.
 * pdf.js may transfer/detach the buffer it receives, so we NEVER
 * pass the original ArrayBuffer directly.
 */
function makePdfJsDataCopy(arrayBuffer: ArrayBuffer): Uint8Array {
  const src = new Uint8Array(arrayBuffer);
  const copy = new Uint8Array(src.byteLength);
  copy.set(src);
  return copy;
}

interface ThumbnailUpdate {
  pageIndex: number;
  dataUrl: string;
}

interface ThumbnailGenerationOptions {
  maxHeight?: number;
  maxWidth?: number;
  concurrency?: number;
  /**
   * Optional subset of 0-based page indices to render.
   * If omitted, all pages are rendered.
   */
  pageIndices?: number[];
  signal?: AbortSignal;
  onThumbnail?: (update: ThumbnailUpdate) => void;
}

/**
 * Unrotated thumbnails – used e.g. in the Split/Extract view.
 */
export async function generateThumbnailsProgressive(
  arrayBuffer: ArrayBuffer,
  options: ThumbnailGenerationOptions = {}
): Promise<string[]> {
  return generateThumbnailsInternal(arrayBuffer, {}, options);
}

/**
 * Thumbnails that respect per-page rotations (for the Reorder view).
 */
export async function generateThumbnailsWithRotationsProgressive(
  arrayBuffer: ArrayBuffer,
  rotations: RotationsMap,
  options: ThumbnailGenerationOptions = {}
): Promise<string[]> {
  return generateThumbnailsInternal(arrayBuffer, rotations, options);
}

async function generateThumbnailsInternal(
  arrayBuffer: ArrayBuffer,
  rotations: RotationsMap,
  options: ThumbnailGenerationOptions = {}
): Promise<string[]> {
  const maxHeight = options.maxHeight ?? 150;
  const maxWidth = options.maxWidth ?? 140;
  const concurrency = Math.max(1, Math.min(options.concurrency ?? 3, 6));
  const signal = options.signal;

  const dataCopy = makePdfJsDataCopy(arrayBuffer);
  const loadingTask = pdfjsLib.getDocument({ data: dataCopy });
  const pdf = await loadingTask.promise;

  const thumbs = Array<string>(pdf.numPages).fill('');

  const pageNums = options.pageIndices
    ? Array.from(
        new Set(
          options.pageIndices
            .filter((pageIndex) => pageIndex >= 0 && pageIndex < pdf.numPages)
            .map((pageIndex) => pageIndex + 1)
        )
      )
    : Array.from({ length: pdf.numPages }, (_, index) => index + 1);

  let nextPageIndex = 0;

  const renderOne = async (pageNum: number) => {
    if (signal?.aborted) return;

    const page = await pdf.getPage(pageNum);
    if (signal?.aborted) return;

    const pageIndex = pageNum - 1;
    const dataUrl = await renderPageThumbnail(
      page,
      pageIndex,
      rotations,
      maxHeight,
      maxWidth
    );

    if (signal?.aborted) return;

    thumbs[pageIndex] = dataUrl;
    options.onThumbnail?.({ pageIndex, dataUrl });
  };

  const worker = async () => {
    while (!signal?.aborted) {
      const pageNum = pageNums[nextPageIndex];
      nextPageIndex += 1;

      if (pageNum == null) return;

      await renderOne(pageNum);

      // Let React/browser paint between batches.
      await new Promise<void>((resolve) => {
        requestAnimationFrame(() => resolve());
      });
    }
  };

  const workerCount = Math.min(concurrency, pageNums.length);
  if (workerCount === 0) return thumbs;

  await Promise.all(Array.from({ length: workerCount }, worker));

  return thumbs;
}

async function renderPageThumbnail(
  page: Awaited<
    ReturnType<
      Awaited<ReturnType<typeof pdfjsLib.getDocument>['promise']>['getPage']
    >
  >,
  originalIndex: number,
  rotations: RotationsMap,
  maxHeight: number,
  maxWidth: number
): Promise<string> {
  const viewport = page.getViewport({ scale: 1 });
  const scaleH = maxHeight / viewport.height;
  const scaleW = maxWidth / viewport.width;
  const scale = Math.min(scaleH, scaleW);
  const scaledViewport = page.getViewport({ scale });

  const baseCanvas = document.createElement('canvas');
  const baseCtx = baseCanvas.getContext('2d');

  if (!baseCtx) return '';

  baseCanvas.width = scaledViewport.width;
  baseCanvas.height = scaledViewport.height;

  const renderTask = page.render({
    canvas: baseCanvas,
    canvasContext: baseCtx,
    viewport: scaledViewport,
  });

  await renderTask.promise;

  const rotationDegRaw = rotations[originalIndex] ?? 0;
  const rotationDeg = ((rotationDegRaw % 360) + 360) % 360;

  if (rotationDeg === 0) {
    return baseCanvas.toDataURL('image/png');
  }

  const rotatedCanvas = document.createElement('canvas');
  const rotatedCtx = rotatedCanvas.getContext('2d');

  if (!rotatedCtx) {
    return baseCanvas.toDataURL('image/png');
  }

  const rad = (rotationDeg * Math.PI) / 180;

  if (rotationDeg === 90 || rotationDeg === 270) {
    rotatedCanvas.width = baseCanvas.height;
    rotatedCanvas.height = baseCanvas.width;
  } else {
    rotatedCanvas.width = baseCanvas.width;
    rotatedCanvas.height = baseCanvas.height;
  }

  rotatedCtx.save();

  switch (rotationDeg) {
    case 90:
      rotatedCtx.translate(rotatedCanvas.width, 0);
      rotatedCtx.rotate(rad);
      break;
    case 180:
      rotatedCtx.translate(rotatedCanvas.width, rotatedCanvas.height);
      rotatedCtx.rotate(rad);
      break;
    case 270:
      rotatedCtx.translate(0, rotatedCanvas.height);
      rotatedCtx.rotate(rad);
      break;
  }

  rotatedCtx.drawImage(baseCanvas, 0, 0);
  rotatedCtx.restore();

  return rotatedCanvas.toDataURL('image/png');
}