PRA-only-demo/packages/rpa-core/src/capabilities/pdf.ts

import { promises as fs } from "node:fs";
import fsSync from "node:fs";
import { fileURLToPath, pathToFileURL } from "node:url";
import { PDFDocument } from "pdf-lib";
import path from "node:path";
import { register } from "../engine/registry";
import { RuntimeError, type CapabilityHandler } from "../engine/types";

function text(args: Record<string, unknown>, keys: string[], fallback = ""): string {
  for (const key of keys) if (args[key] !== undefined && args[key] !== null && args[key] !== "") return String(args[key]);
  return fallback;
}

function required(value: string, label: string, nodeId: string, componentId: string): string {
  if (!value.trim()) throw new RuntimeError("validation_error", `${label}为空`, { nodeId, componentId });
  return path.resolve(value);
}

const pageCount: CapabilityHandler = async (_ctx, node, args) => {
  const filePath = required(text(args, ["fileInPath", "path", "file"]), "PDF文件路径", node.id, node.componentId);
  const document = await PDFDocument.load(await fs.readFile(filePath), { ignoreEncryption: true });
  return document.getPageCount();
};

const concat: CapabilityHandler = async (_ctx, node, args) => {
  const folder = required(text(args, ["newPdfPath", "folder", "path"]), "输出目录", node.id, node.componentId);
  const name = text(args, ["newPdfName", "name"], "合并结果.pdf");
  const output = path.join(folder, path.extname(name) ? name : `${name}.pdf`);
  const rawList = text(args, ["pdfList", "files"]);
  let files: string[] = [];
  try {
    const parsed = JSON.parse(rawList);
    files = Array.isArray(parsed) ? parsed.map(String) : [String(parsed)];
  } catch {
    files = rawList.split(/[\r\n;,]+/).map((item) => item.trim()).filter(Boolean);
  }
  if (!files.length) throw new RuntimeError("validation_error", "待合并 PDF 列表为空", { nodeId: node.id, componentId: node.componentId });
  const target = await PDFDocument.create();
  for (const file of files) {
    const source = await PDFDocument.load(await fs.readFile(path.resolve(file)), { ignoreEncryption: true });
    const pages = await target.copyPages(source, source.getPageIndices());
    pages.forEach((page) => target.addPage(page));
  }
  await fs.mkdir(path.dirname(output), { recursive: true });
  await fs.writeFile(output, await target.save());
  return output;
};

register(["PDF_GetPdfPageNum", "PDF_GetPageCount"], pageCount);
register("PDF_ConcatPdf", concat);
register("PDF_GetPdfText", async (_ctx, node, args) => {
  const filePath = required(text(args, ["fileInPath", "path", "file"]), "PDF文件路径", node.id, node.componentId);
  const pdfjs = await import("pdfjs-dist/legacy/build/pdf.mjs");
  const document = await pdfjs.getDocument({ data: new Uint8Array(await fs.readFile(filePath)) }).promise;
  const pages = text(args, ["pdfNum", "page"]);
  const indices = pages ? pages.split(/[ ,，;；]+/).map((item) => Number(item)).filter((item) => item > 0) : Array.from({ length: document.numPages }, (_, index) => index + 1);
  const chunks: string[] = [];
  for (const index of indices) {
    const page = await document.getPage(index);
    const content = await page.getTextContent();
    chunks.push(content.items.map((item) => ("str" in item ? item.str : "")).join(""));
  }
  return chunks.join("\n");
});

type PdfJsDocument = Awaited<ReturnType<typeof import("pdfjs-dist/legacy/build/pdf.mjs")["getDocument"]>["promise"]>;

function standardFontUrl(): string | undefined {
  const candidates = [
    process.env.RPA_PDFJS_STANDARD_FONTS,
    path.join(process.cwd(), "node_modules", "pdfjs-dist", "standard_fonts"),
    path.join(process.execPath, "resources", "app.asar", "node_modules", "pdfjs-dist", "standard_fonts"),
    path.join(process.execPath, "..", "resources", "app.asar", "node_modules", "pdfjs-dist", "standard_fonts")
  ].filter(Boolean) as string[];
  const found = candidates.find((candidate) => fsSync.existsSync(candidate));
  return found ? pathToFileURL(path.join(found, path.sep)).href : undefined;
}

class LocalStandardFontFactory {
  baseUrl: string;
  constructor(options: { baseUrl: string }) { this.baseUrl = options.baseUrl; }
  async fetch({ filename }: { filename: string }): Promise<Uint8Array> {
    return new Uint8Array(await fs.readFile(path.join(fileURLToPath(this.baseUrl), filename)));
  }
}

async function loadPdf(filePath: string, password: string) {
  const canvasModule = await import("@napi-rs/canvas");
  const globals = globalThis as unknown as Record<string, unknown>;
  globals.Path2D = canvasModule.Path2D;
  globals.DOMMatrix = canvasModule.DOMMatrix;
  globals.ImageData = canvasModule.ImageData;
  const pdfjs = await import("pdfjs-dist/legacy/build/pdf.mjs");
  return { pdfjs, document: await pdfjs.getDocument({ data: new Uint8Array(await fs.readFile(filePath)), password: password || undefined, isEvalSupported: false, useSystemFonts: false, standardFontDataUrl: standardFontUrl(), StandardFontDataFactory: LocalStandardFontFactory }).promise };
}

function pageIndices(document: PdfJsDocument, raw: unknown): number[] {
  const range = String(raw ?? "all").trim().toLowerCase();
  if (!range || range === "all") return Array.from({ length: document.numPages }, (_, index) => index + 1);
  const result = new Set<number>();
  for (const part of range.split(/[,，;；\s]+/).filter(Boolean)) {
    const match = part.match(/^(\d+)-(\d+)$/);
    if (match) {
      const start = Math.max(1, Number(match[1])); const end = Math.min(document.numPages, Number(match[2]));
      for (let index = start; index <= end; index += 1) result.add(index);
    } else {
      const index = Number(part);
      if (index >= 1 && index <= document.numPages) result.add(Math.floor(index));
    }
  }
  const values = [...result].sort((left, right) => left - right);
  if (!values.length) throw new Error(`PDF 页码范围无效：${range}`);
  return values;
}

async function renderPage(page: Awaited<ReturnType<PdfJsDocument["getPage"]>>, scale = 2): Promise<Buffer> {
  const canvasModule = await import("@napi-rs/canvas");
  const viewport = page.getViewport({ scale });
  const canvas = canvasModule.createCanvas(Math.ceil(viewport.width), Math.ceil(viewport.height));
  const context = canvas.getContext("2d");
  context.fillStyle = "white"; context.fillRect(0, 0, canvas.width, canvas.height);
  await page.render({ canvasContext: context as unknown as CanvasRenderingContext2D, viewport }).promise;
  return canvas.toBuffer("image/png");
}

function saveImage(folder: string, prefix: string, suffix: string, data: Buffer, type: string): string {
  const ext = /jpe?g/i.test(type) ? "jpg" : "png";
  const output = path.join(path.resolve(folder), `${prefix}_${suffix}_${Date.now()}.${ext}`);
  fsSync.mkdirSync(path.dirname(output), { recursive: true });
  fsSync.writeFileSync(output, data);
  return output;
}

async function rawImageToBuffer(image: Record<string, unknown>, type: string): Promise<Buffer> {
  const width = Number(image.width || 0); const height = Number(image.height || 0);
  const data = image.data;
  if (!width || !height || !(data instanceof Uint8Array) || data.length === 0) throw new Error("PDF 图片数据无效");
  if (data[0] === 0xff && data[1] === 0xd8) return Buffer.from(data);
  const canvasModule = await import("@napi-rs/canvas");
  const canvas = canvasModule.createCanvas(width, height); const context = canvas.getContext("2d");
  const target = context.createImageData(width, height); const kind = Number(image.kind || 0);
  if (kind === 1) {
    const bytesPerRow = Math.ceil(width / 8);
    for (let y = 0; y < height; y += 1) for (let x = 0; x < width; x += 1) {
      const bit = (data[y * bytesPerRow + (x >> 3)] >> (7 - (x & 7))) & 1; const output = (y * width + x) * 4; const color = bit ? 255 : 0;
      target.data[output] = color; target.data[output + 1] = color; target.data[output + 2] = color; target.data[output + 3] = 255;
    }
  } else if (kind === 3) { target.data.set(data.subarray(0, Math.min(data.length, target.data.length))); }
  else {
    for (let source = 0, targetIndex = 0; source < data.length && targetIndex < target.data.length; source += 3, targetIndex += 4) {
      target.data[targetIndex] = data[source]; target.data[targetIndex + 1] = data[source + 1] ?? 0; target.data[targetIndex + 2] = data[source + 2] ?? 0; target.data[targetIndex + 3] = 255;
    }
  }
  context.putImageData(target, 0, 0);
  return /jpe?g/i.test(type) ? canvas.toBuffer("image/jpeg") : canvas.toBuffer("image/png");
}

async function objectImage(page: Awaited<ReturnType<PdfJsDocument["getPage"]>>, objId: unknown): Promise<Record<string, unknown> | null> {
  if (typeof objId !== "string") return null;
  const objects = objId.startsWith("g_") ? page.commonObjs : page.objs;
  try { if (objects.has(objId)) return objects.get(objId) as Record<string, unknown>; } catch {}
  return await new Promise((resolve) => { try { objects.get(objId, (value: unknown) => resolve(value as Record<string, unknown>)); } catch { resolve(null); } });
}

const encryption: CapabilityHandler = async (_ctx, node, args) => {
  const filePath = required(text(args, ["fileInPath", "path", "file"]), "PDF文件路径", node.id, node.componentId);
  const password = String(args.passWord ?? args.password ?? "");
  if (!password) throw new RuntimeError("validation_error", "PDF密码为空", { nodeId: node.id, componentId: node.componentId });
  const { PDFDocument } = await import("@cantoo/pdf-lib");
  const document = await PDFDocument.load(await fs.readFile(filePath), { ignoreEncryption: true });
  document.encrypt({ userPassword: password, ownerPassword: password, permissions: { printing: "highResolution", copying: true, modifying: false } });
  const temporary = `${filePath}.${Date.now()}.enc`;
  await fs.writeFile(temporary, await document.save());
  await fs.rename(temporary, filePath);
  return filePath;
};

const extractImages: CapabilityHandler = async (_ctx, node, args) => {
  const filePath = required(text(args, ["fileInPath", "path", "file"]), "PDF文件路径", node.id, node.componentId);
  const folder = required(text(args, ["imgSavePath", "pdfSavePath", "folder", "path"]), "图片保存路径", node.id, node.componentId);
  const type = text(args, ["imageType", "type"], "png");
  const prefix = text(args, ["imageName", "name"], "pdf-image");
  const { pdfjs, document } = await loadPdf(filePath, String(args.pdfPassWord ?? args.password ?? ""));
  const outputs: string[] = [];
  for (const pageNumber of pageIndices(document, args.pdfNum)) {
    const page = await document.getPage(pageNumber);
    const opList = await page.getOperatorList();
    let sequence = 0;
    for (let index = 0; index < opList.fnArray.length; index += 1) {
      const op = opList.fnArray[index]; const commandArgs = opList.argsArray[index] as unknown[];
      let image: Record<string, unknown> | null = null;
      if (op === pdfjs.OPS.paintImageXObject || op === pdfjs.OPS.paintImageXObjectRepeat) image = await objectImage(page, commandArgs[0]);
      else if (op === pdfjs.OPS.paintInlineImageXObject) image = commandArgs[0] as Record<string, unknown>;
      else if (op === pdfjs.OPS.paintInlineImageXObjectGroup) image = commandArgs[0] as Record<string, unknown>;
      if (!image) continue;
      sequence += 1;
      outputs.push(saveImage(folder, `${prefix}-p${pageNumber}`, String(sequence), await rawImageToBuffer(image, type), type));
    }
  }
  return outputs;
};

const pdfToImages: CapabilityHandler = async (_ctx, node, args) => {
  const filePath = required(text(args, ["fileInPath", "path", "file"]), "PDF文件路径", node.id, node.componentId);
  const folder = required(text(args, ["pdfSavePath", "imgSavePath", "folder", "path"]), "图片保存路径", node.id, node.componentId);
  const type = text(args, ["imageType", "type"], "png"); const prefix = text(args, ["imageName", "name"], "pdf-page");
  const { document } = await loadPdf(filePath, String(args.pdfPassWord ?? args.password ?? ""));
  const outputs: string[] = [];
  for (const pageNumber of pageIndices(document, args.pdfNum)) {
    const buffer = await renderPage(await document.getPage(pageNumber));
    outputs.push(saveImage(folder, prefix, String(pageNumber), buffer, type));
  }
  return outputs;
};

interface TextToken { x: number; y: number; value: string; }

async function extractTableRows(page: Awaited<ReturnType<PdfJsDocument["getPage"]>>): Promise<string[][]> {
  const content = await page.getTextContent();
    const tokens = content.items.flatMap((item) => {
      if (!("str" in item) || !item.str.trim()) return [];
      const transform = (item as { transform?: number[] }).transform ?? [1,0,0,1,0,0];
      return [{ x: Number(transform[4] || 0), y: Number(transform[5] || 0), value: item.str.trim() }];
    });
    if (tokens.length < 2) return [];
    const lines = new Map<number, TextToken[]>();
    for (const token of [...tokens].sort((left, right) => right.y - left.y || left.x - right.x)) {
      const lineKey = [...lines.keys()].find((key) => Math.abs(key - token.y) <= 3.5) ?? token.y;
      if (!lines.has(lineKey)) lines.set(lineKey, []);
      lines.get(lineKey)!.push(token);
    }
    const lineList = [...lines.entries()].map(([y, values]) => ({ y, values: values.sort((left, right) => left.x - right.x) })).sort((left, right) => right.y - left.y);
    const counts = new Map<number, number>();
    for (const line of lineList) for (const token of line.values) {
      const key = Math.round(token.x / 5) * 5;
      counts.set(key, (counts.get(key) || 0) + 1);
    }
    const columns = [...counts.entries()].filter(([, count]) => count >= Math.min(2, lineList.length)).sort((left, right) => left[0] - right[0]).map(([key]) => key);
    if (columns.length < 2) return [];
    const rows: string[][] = [];
    for (const line of lineList) {
      const cells = columns.map(() => "");
      for (const token of line.values) {
        let nearest = 0;
        for (let index = 1; index < columns.length; index += 1) if (Math.abs(token.x - columns[index]) < Math.abs(token.x - columns[nearest])) nearest = index;
        cells[nearest] = cells[nearest] ? `${cells[nearest]} ${token.value}` : token.value;
      }
      if (cells.filter(Boolean).length >= 2) rows.push(cells);
    }
    return rows.length >= 2 ? rows : [];
}

const readTables: CapabilityHandler = async (_ctx, node, args) => {
  const filePath = required(text(args, ["fileInPath", "path", "file"]), "PDF文件路径", node.id, node.componentId);
  const { document } = await loadPdf(filePath, String(args.pdfPassWord ?? args.password ?? ""));
  const rows: string[][] = [];
  for (const pageNumber of pageIndices(document, args.pdfNum)) rows.push(...await extractTableRows(await document.getPage(pageNumber)));
  return rows;
};

register("PDF_PdfEncryption", encryption);
register("PDF_GetPdfImg", extractImages);
register("PDF_PdfToImg", pdfToImages);
register("PDF_GetPdfTable", readTables);