pdf-tools/src/pdf/pdfService.ts

import { PDFDocument, degrees } from 'pdf-lib';
import type { PdfFile, PageRef, SplitResult, Range } from './pdfTypes';

function createId() {
  return Math.random().toString(36).slice(2);
}

function pdfBytesToArrayBuffer(bytes: Uint8Array): ArrayBuffer {
  const buffer = new ArrayBuffer(bytes.byteLength);
  new Uint8Array(buffer).set(bytes);
  return buffer;
}

function pdfBytesToBlob(bytes: Uint8Array): Blob {
  return new Blob([pdfBytesToArrayBuffer(bytes)], { type: 'application/pdf' });
}

export async function loadPdfFromFile(file: File): Promise<PdfFile> {
  const arrayBuffer = await file.arrayBuffer();
  const doc = await PDFDocument.load(arrayBuffer);

  return {
    id: createId(),
    name: file.name,
    doc,
    pageCount: doc.getPageCount(),
    arrayBuffer,
  };
}

export async function mergePdfFiles(
  basePdf: PdfFile,
  newPdf: PdfFile,
  insertAt: number
): Promise<PdfFile> {
  const baseDoc = basePdf.doc ?? (await PDFDocument.load(basePdf.arrayBuffer));
  const newDoc = newPdf.doc ?? (await PDFDocument.load(newPdf.arrayBuffer));

  const mergedDoc = await PDFDocument.create();

  const basePageCount = baseDoc.getPageCount();
  const newPageCount = newDoc.getPageCount();

  const clampedInsertAt = Math.min(Math.max(insertAt, 0), basePageCount);

  const basePages = await mergedDoc.copyPages(
    baseDoc,
    Array.from({ length: basePageCount }, (_, i) => i)
  );
  const newPages = await mergedDoc.copyPages(
    newDoc,
    Array.from({ length: newPageCount }, (_, i) => i)
  );

  for (let i = 0; i < clampedInsertAt; i += 1) {
    mergedDoc.addPage(basePages[i]);
  }
  for (let i = 0; i < newPages.length; i += 1) {
    mergedDoc.addPage(newPages[i]);
  }
  for (let i = clampedInsertAt; i < basePages.length; i += 1) {
    mergedDoc.addPage(basePages[i]);
  }

  const bytes = await mergedDoc.save();
  const buffer = pdfBytesToArrayBuffer(bytes);

  const baseName = basePdf.name.replace(/\.pdf$/i, '');
  const newName = newPdf.name.replace(/\.pdf$/i, '');

  return {
    id: createId(),
    name: `${baseName}_plus_${newName}.pdf`,
    arrayBuffer: buffer,
    pageCount: mergedDoc.getPageCount(),
    doc: mergedDoc,
  };
}

interface MergePdfFilesAtPositionOptions {
  basePdf: PdfFile | null;
  incomingPdfs: PdfFile[];
  insertAt: number;
  name: string;
}

export async function mergePdfFilesAtPosition({
  basePdf,
  incomingPdfs,
  insertAt,
  name,
}: MergePdfFilesAtPositionOptions): Promise<PdfFile> {
  if (!basePdf && incomingPdfs.length === 0) {
    throw new Error('At least one PDF is required for merging');
  }

  const mergedDoc = await PDFDocument.create();

  const addAllPages = async (sourcePdf: PdfFile) => {
    const sourceDoc =
      sourcePdf.doc ?? (await PDFDocument.load(sourcePdf.arrayBuffer));
    const pageCount = sourceDoc.getPageCount();
    const pages = await mergedDoc.copyPages(
      sourceDoc,
      Array.from({ length: pageCount }, (_, i) => i)
    );
    pages.forEach((page) => mergedDoc.addPage(page));
  };

  if (!basePdf) {
    for (const incomingPdf of incomingPdfs) {
      await addAllPages(incomingPdf);
    }
  } else {
    const baseDoc =
      basePdf.doc ?? (await PDFDocument.load(basePdf.arrayBuffer));
    const basePageCount = baseDoc.getPageCount();
    const clampedInsertAt = Math.min(Math.max(insertAt, 0), basePageCount);

    const basePages = await mergedDoc.copyPages(
      baseDoc,
      Array.from({ length: basePageCount }, (_, i) => i)
    );

    for (let i = 0; i < clampedInsertAt; i += 1) {
      mergedDoc.addPage(basePages[i]);
    }

    for (const incomingPdf of incomingPdfs) {
      await addAllPages(incomingPdf);
    }

    for (let i = clampedInsertAt; i < basePages.length; i += 1) {
      mergedDoc.addPage(basePages[i]);
    }
  }

  const bytes = await mergedDoc.save();
  const buffer = pdfBytesToArrayBuffer(bytes);

  return {
    id: createId(),
    name,
    arrayBuffer: buffer,
    pageCount: mergedDoc.getPageCount(),
    doc: mergedDoc,
  };
}

export async function splitIntoSinglePages(
  pdf: PdfFile
): Promise<SplitResult[]> {
  const { doc, name } = pdf;

  const title = doc.getTitle();
  const author = doc.getAuthor();
  const subject = doc.getSubject();
  const keywords = doc.getKeywords();
  const producer = doc.getProducer();
  const creator = doc.getCreator();
  const creationDate = doc.getCreationDate();
  const modificationDate = doc.getModificationDate();

  const results: SplitResult[] = [];

  for (let i = 0; i < doc.getPageCount(); i++) {
    const newDoc = await PDFDocument.create();
    const [copiedPage] = await newDoc.copyPages(doc, [i]);
    newDoc.addPage(copiedPage);

    if (title) newDoc.setTitle(title);
    if (author) newDoc.setAuthor(author);
    if (subject) newDoc.setSubject(subject);
    if (keywords) newDoc.setKeywords([keywords]);
    if (producer) newDoc.setProducer(producer);
    if (creator) newDoc.setCreator(creator);
    if (creationDate) newDoc.setCreationDate(creationDate);
    if (modificationDate) newDoc.setModificationDate(modificationDate);

    const bytes = await newDoc.save();
    const blob = pdfBytesToBlob(bytes);

    const base = name.replace(/\.pdf$/i, '');
    const filename = `${base}_page_${String(i + 1).padStart(3, '0')}.pdf`;

    results.push({
      pageIndex: i,
      blob,
      filename,
    });
  }

  return results;
}

export async function extractRange(pdf: PdfFile, range: Range): Promise<Blob> {
  const { doc } = pdf;
  const pageCount = doc.getPageCount();

  const fromIndex = Math.max(0, range.from - 1);
  const toIndex = Math.min(pageCount - 1, range.to - 1);

  if (fromIndex > toIndex) {
    throw new Error('Invalid range: from > to');
  }

  const newDoc = await PDFDocument.create();
  const indices: number[] = [];
  for (let i = fromIndex; i <= toIndex; i++) indices.push(i);

  const copiedPages = await newDoc.copyPages(doc, indices);
  copiedPages.forEach((p) => newDoc.addPage(p));

  const bytes = await newDoc.save();
  return pdfBytesToBlob(bytes);
}

export async function mergePdfs(pdfs: PdfFile[]): Promise<Blob> {
  const newDoc = await PDFDocument.create();

  for (const pdf of pdfs) {
    const pageCount = pdf.doc.getPageCount();
    const indices = Array.from({ length: pageCount }, (_, i) => i);
    const copiedPages = await newDoc.copyPages(pdf.doc, indices);
    copiedPages.forEach((p) => newDoc.addPage(p));
  }

  const bytes = await newDoc.save();
  return pdfBytesToBlob(bytes);
}

export async function exportPages(
  pdf: PdfFile,
  pages: PageRef[]
): Promise<Blob> {
  const { doc } = pdf;
  const pageCount = doc.getPageCount();

  if (pages.length === 0) {
    throw new Error('Pages must contain at least one page');
  }

  if (
    pages.some(
      (page) => page.sourcePageIndex < 0 || page.sourcePageIndex >= pageCount
    )
  ) {
    throw new Error('Pages contain invalid source page indices');
  }

  const newDoc = await PDFDocument.create();
  const indices = pages.map((page) => page.sourcePageIndex);

  const copiedPages = await newDoc.copyPages(doc, indices);

  copiedPages.forEach((page, idx) => {
    const angle = pages[idx].rotation;

    if (typeof angle === 'number' && angle % 360 !== 0) {
      page.setRotation(degrees(angle));
    }

    newDoc.addPage(page);
  });

  const bytes = await newDoc.save();
  return pdfBytesToBlob(bytes);
}

export async function exportReordered(
  pdf: PdfFile,
  order: number[],
  rotations?: Record<number, number>
): Promise<Blob> {
  return exportPages(
    pdf,
    order.map((sourcePageIndex) => ({
      id: String(sourcePageIndex),
      sourcePageIndex,
      rotation: rotations?.[sourcePageIndex] ?? 0,
    }))
  );
}