import { PDFDocument, PDFName, decodePDFRawStream, PDFRawStream } from 'pdf-lib';
import fs from 'fs';
import path from 'path';

async function run() {
  const inputPath = process.argv[2];
  const outputPath = process.argv[3] || inputPath;

  if (!inputPath || !fs.existsSync(inputPath)) {
    console.error(`[ERROR] Input file not found: ${inputPath}`);
    process.exit(1);
  }

  console.log(`[INFO] Injecting PANTONE 289 C into: ${path.basename(inputPath)}`);
  const pdfBytes = fs.readFileSync(inputPath);
  const pdfDoc = await PDFDocument.load(pdfBytes);

  // CMYK tint transform for PANTONE 289 C.
  // Approximation: 100C, 64M, 0Y, 60K (1.0 0.64 0.0 0.60)
  // The function maps an input x (0 to 1) to CMYK: x*1.0, x*0.64, x*0.0, x*0.60
  const tintTransform = pdfDoc.context.obj({
    FunctionType: 2,
    Domain: [0, 1],
    C0: [0, 0, 0, 0],
    C1: [1.0, 0.64, 0.0, 0.6],
    N: 1,
  });

  // Create the /Separation color space array
  const pantoneCS = pdfDoc.context.obj([
    PDFName.of('Separation'),
    PDFName.of('PANTONE#20289#20C'),
    PDFName.of('DeviceCMYK'),
    tintTransform,
  ]);
  const csRef = pdfDoc.context.register(pantoneCS);

  const pages = pdfDoc.getPages();
  for (const page of pages) {
    let resources = page.node.Resources();
    if (!resources) {
      resources = pdfDoc.context.obj({});
      page.node.set(PDFName.of('Resources'), resources);
    }

    // Ensure ColorSpace dict exists
    let colorSpaces = resources.get(PDFName.of('ColorSpace')) as any;
    if (!colorSpaces) {
      colorSpaces = pdfDoc.context.obj({});
      resources.set(PDFName.of('ColorSpace'), colorSpaces);
    }
    // Register the PANTONE color space as /CSPantone in the page's resources
    colorSpaces.set(PDFName.of('CSPantone'), csRef);

    // Read all content streams
    const contents = page.node.Contents();
    let streams: any[] = [];
    if (contents) {
      const resolvedContents = pdfDoc.context.lookup(contents);
      if (resolvedContents instanceof PDFRawStream) {
        streams = [resolvedContents];
      } else if (resolvedContents.constructor.name === 'PDFArray') {
        const arr = resolvedContents as any;
        for (let j = 0; j < arr.size(); j++) {
          streams.push(pdfDoc.context.lookup(arr.get(j)));
        }
      }
    }

    let newStr = '';
    for (let i = 0; i < streams.length; i++) {
      const stream = streams[i];
      if (stream instanceof PDFRawStream) {
        const decoded = decodePDFRawStream(stream).decode();
        const str = Buffer.from(decoded).toString('utf-8');
        newStr += str + '\n';
      }
    }

    // The CMYK values for the KLZ Blue created by Ghostscript from #000a66
    // is '1 0.98 0.224 0.298'. We use a slightly dynamic regex to catch small precision diffs.
    const dynamicRegexK = /1 0\.98[0-9]* 0\.224[0-9]* 0\.298[0-9]* k/g;
    const dynamicRegexStrokeK = /1 0\.98[0-9]* 0\.224[0-9]* 0\.298[0-9]* K/g;

    let modified = false;
    if (dynamicRegexK.test(newStr) || dynamicRegexStrokeK.test(newStr)) {
      newStr = newStr.replace(dynamicRegexK, '/CSPantone cs 1 scn');
      newStr = newStr.replace(dynamicRegexStrokeK, '/CSPantone CS 1 SCN');
      modified = true;
    }

    if (modified) {
      const newStream = pdfDoc.context.flateStream(newStr);
      const newRef = pdfDoc.context.register(newStream);
      page.node.set(PDFName.of('Contents'), newRef);
      console.log(`[INFO] Spot color injected into page ${pages.indexOf(page) + 1}.`);
    }
  }

  const finalPdfBytes = await pdfDoc.save();
  fs.writeFileSync(outputPath, finalPdfBytes);
  console.log(`[SUCCESS] Saved 5-color PDF to ${outputPath}`);
}

run().catch(console.error);