generate-caption-pngs-word-efficient.js

#!/usr/bin/env node

const fs = require("fs");
const path = require("path");
const { createCanvas, registerFont } = require("canvas");
const mkdirp = require("mkdirp");

// Configuration
const FPS = 30;
const WIDTH = 1280;
const HEIGHT = 720;
const CAPTION_AREA_HEIGHT = 400;
const PADDING_X = 15;
const PADDING_Y = 5;
const BG_COLOR = "rgba(18, 46, 73, 1)";
const ACTIVE_COLOR = "#fcfaef";
const INACTIVE_COLOR = "#bbbbbb";
const ACTIVE_HIGHLIGHT_COLOR = "#76b1af"; // Bright green highlight color
const INACTIVE_HIGHLIGHT_COLOR = "#547e7c"; // Dark green highlight color
const BORDER_RADIUS = 20;
const REMOVE_PERIODS = true;
const TEXT_CASE = "lowercase";
const FONT_SIZE = 60;
const LINE_HEIGHT = 75;
const UNSPOKEN_OPACITY = 0.0; // Opacity for unspoken words (0 = invisible, 1 = fully visible)

// Drop shadow configuration
const SHADOW_COLOR = "rgba(0, 0, 0, 0.6)";
const INACTIVE_SHADOW_BLUR = 2;
const INACTIVE_SHADOW_OFFSET_X = 2;
const INACTIVE_SHADOW_OFFSET_Y = 2;
const ACTIVE_SHADOW_BLUR = 4;
const ACTIVE_SHADOW_OFFSET_X = 3;
const ACTIVE_SHADOW_OFFSET_Y = 3;

// Parse command line arguments
const args = process.argv.slice(2);
let jsonFile = null;
let outputDir = null;
let startTime = 0;
let endTime = Infinity;
let removePeriods = REMOVE_PERIODS;
let textCase = TEXT_CASE;
let hideUnspoken = false;
let generateFirstWordOnly = false;

for (let i = 0; i < args.length; i++) {
  if (args[i] === "-json" && i + 1 < args.length) {
    jsonFile = args[i + 1];
    i++;
  } else if (args[i] === "-output" && i + 1 < args.length) {
    outputDir = args[i + 1];
    i++;
  } else if (args[i] === "-start" && i + 1 < args.length) {
    startTime = parseFloat(args[i + 1]);
    i++;
  } else if (args[i] === "-end" && i + 1 < args.length) {
    endTime = parseFloat(args[i + 1]);
    i++;
  } else if (args[i] === "-remove-periods" && i + 1 < args.length) {
    removePeriods = args[i + 1].toLowerCase() === "true";
    i++;
  } else if (args[i] === "-case" && i + 1 < args.length) {
    textCase = args[i + 1].toLowerCase();
    if (!["lowercase", "uppercase", "original"].includes(textCase)) {
      console.warn(
        `Invalid case option "${textCase}", using default "${TEXT_CASE}"`
      );
      textCase = TEXT_CASE;
    }
    i++;
  } else if (args[i] === "-hide-unspoken" && i + 1 < args.length) {
    hideUnspoken = args[i + 1].toLowerCase() === "true";
    console.log(`Hide unspoken words: ${hideUnspoken}`);
    i++;
  } else if (args[i] === "-generate-first-word-only" && i + 1 < args.length) {
    generateFirstWordOnly = args[i + 1].toLowerCase() === "true";
    i++;
  }
}

if (!jsonFile) {
  console.error(
    "Usage: node generate-caption-pngs-word-efficient.js -json [group.json] -output [output_dir] -start [start_time] -end [end_time] -remove-periods [true/false] -case [lowercase/uppercase/original] -hide-unspoken [true/false] -generate-first-word-only [true/false]"
  );
  process.exit(1);
}

if (!outputDir) {
  // Default output directory
  const jsonBaseName = path.basename(jsonFile, path.extname(jsonFile));
  outputDir = `${jsonBaseName}_word_frames`;
}

// Create output directory if it doesn't exist
try {
  mkdirp.sync(outputDir);
} catch (error) {
  console.error(`Error creating output directory: ${error.message}`);
  process.exit(1);
}

// Try to register fonts - adjust path as needed for your system
try {
  // Different font paths for different operating systems
  if (process.platform === "win32") {
    registerFont("C:\\Windows\\Fonts\\arial.ttf", { family: "Arial" });
    registerFont("C:\\Windows\\Fonts\\arialbd.ttf", {
      family: "Arial",
      weight: "bold",
    });
    try {
      registerFont("A:\\Procaci-Amamenta\\OUTPUT\\Reels\\Metropolis-Bold.ttf", {
        family: "Metropolis",
        weight: "bold",
      });
    } catch (error) {
      console.warn(
        `Warning: Could not register Metropolis font: ${error.message}`
      );
    }
  } else if (process.platform === "darwin") {
    registerFont("/Library/Fonts/Arial.ttf", { family: "Arial" });
    registerFont("/Library/Fonts/Arial Bold.ttf", {
      family: "Arial",
      weight: "bold",
    });
  } else {
    registerFont("/usr/share/fonts/truetype/msttcorefonts/Arial.ttf", {
      family: "Arial",
    });
    registerFont("/usr/share/fonts/truetype/msttcorefonts/Arial_Bold.ttf", {
      family: "Arial",
      weight: "bold",
    });
  }
} catch (error) {
  console.warn(`Warning: Could not register Arial font: ${error.message}`);
  console.warn("Using system default font instead");
}

// Read the JSON file
console.log(`Reading JSON file: ${jsonFile}`);
const data = fs.readFileSync(jsonFile, "utf8");
const captionsData = JSON.parse(data);

// Calculate end time if not specified
if (endTime === Infinity) {
  endTime = Math.max(...captionsData.captions.map((caption) => caption.end));
}

console.log(
  `Processing captions from ${startTime}s to ${endTime}s, found ${captionsData.captions.length} captions`
);

// Function to format text based on options
function formatText(text, forceUpperCase = false) {
  let formatted = text;

  // Remove periods if enabled
  if (removePeriods) {
    formatted = formatted.replace(/\./g, "");
  }

  // Apply case formatting
  if (forceUpperCase) {
    formatted = formatted.toUpperCase();
  } else if (textCase === "lowercase") {
    formatted = formatted.toLowerCase();
  } else if (textCase === "uppercase") {
    formatted = formatted.toUpperCase();
  }

  return formatted;
}

// Helper function to draw a rounded rectangle
function drawRoundedRect(ctx, x, y, width, height, radius) {
  ctx.beginPath();
  ctx.moveTo(x + radius, y);
  ctx.lineTo(x + width - radius, y);
  ctx.quadraticCurveTo(x + width, y, x + width, y + radius);
  ctx.lineTo(x + width, y + height - radius);
  ctx.quadraticCurveTo(x + width, y + height, x + width - radius, y + height);
  ctx.lineTo(x + radius, y + height);
  ctx.quadraticCurveTo(x, y + height, x, y + height - radius);
  ctx.lineTo(x, y + radius);
  ctx.quadraticCurveTo(x, y, x + radius, y);
  ctx.closePath();
  ctx.fill();
}

// Function to split caption words into lines for optimal balance
function splitCaptionIntoLines(
  captionWordsData,
  ctx,
  maxWidth,
  idealMaxCharsSingleLine,
  formatTextFunc
) {
  const lines = []; // Array of arrays of word objects
  if (!captionWordsData || captionWordsData.length === 0) return lines;

  // captionWordsData is an array of {word: string, originalWordIndex: number, lineBreak?: boolean}
  const allWordsWithFormatting = captionWordsData.map((wd) => ({
    word: formatTextFunc(wd.word.trim()), // Apply formatting
    originalIndex: wd.originalWordIndex,
    lineBreakHint: wd.lineBreak,
    highlighted: wd.highlighted,
    upperCase: wd.upperCase,
  }));

  const fullCaptionText = allWordsWithFormatting.map((w) => w.word).join(" ");
  const fullCaptionWidth = ctx.measureText(fullCaptionText).width;

  // 1. Check for explicit line breaks from data
  let explicitBreakFound = false;
  for (let i = 0; i < allWordsWithFormatting.length - 1; i++) {
    if (allWordsWithFormatting[i].lineBreakHint) {
      const line1Words = allWordsWithFormatting.slice(0, i + 1);
      const line2Words = allWordsWithFormatting.slice(i + 1);
      if (line1Words.length > 0) lines.push(line1Words);
      if (line2Words.length > 0) lines.push(line2Words);
      explicitBreakFound = true;
      break;
    }
  }
  if (explicitBreakFound) {
    // If more than 2 lines due to multiple hints, consolidate to 2 for now.
    if (lines.length > 2) {
      console.warn(
        "Multiple explicit line breaks resulted in >2 lines. Consolidating."
      );
      const combined = [].concat(...lines.slice(1));
      lines.splice(1, lines.length - 1);
      if (combined.length > 0) lines.push(combined);
    }
    return lines.filter((line) => line.length > 0);
  }

  // 2. If it fits on one line (both width and a reasonable character count for a single line)
  if (
    fullCaptionWidth <= maxWidth &&
    fullCaptionText.length <= idealMaxCharsSingleLine
  ) {
    lines.push(allWordsWithFormatting);
    return lines;
  }

  // 3. Needs to be split into two lines: find the best split point
  let bestSplit = {
    index: -1, // Index of the word *before* which the second line starts
    diff: Infinity,
    line1Words: [],
    line2Words: [],
  };

  // Iterate through possible split points (split occurs AFTER word i)
  for (let i = 0; i < allWordsWithFormatting.length - 1; i++) {
    const currentLine1Words = allWordsWithFormatting.slice(0, i + 1);
    const currentLine2Words = allWordsWithFormatting.slice(i + 1);

    if (currentLine1Words.length === 0 || currentLine2Words.length === 0)
      continue;

    const line1Text = currentLine1Words.map((w) => w.word).join(" ");
    const line2Text = currentLine2Words.map((w) => w.word).join(" ");

    const line1Width = ctx.measureText(line1Text).width;
    const line2Width = ctx.measureText(line2Text).width;

    // Constraint: Neither line should exceed maxWidth
    if (line1Width > maxWidth || line2Width > maxWidth) {
      continue;
    }

    const diff = Math.abs(line1Text.length - line2Text.length);

    if (diff < bestSplit.diff) {
      bestSplit = {
        index: i, // Word i is the last on the first line
        diff: diff,
        line1Words: currentLine1Words,
        line2Words: currentLine2Words,
      };
    }
  }

  if (bestSplit.index !== -1) {
    if (bestSplit.line1Words.length > 0) lines.push(bestSplit.line1Words);
    if (bestSplit.line2Words.length > 0) lines.push(bestSplit.line2Words);
  } else {
    // Fallback: No balanced split found where both lines are within maxWidth.
    // Use a greedy approach: fill first line as much as possible within maxWidth.
    console.warn(
      `Could not find an ideal balanced two-line split for: "${fullCaptionText.substring(
        0,
        50
      )}...". Using greedy wrap.`
    );
    lines.length = 0;
    let currentLineAggregator = [];
    for (const wordObj of allWordsWithFormatting) {
      const testLineAggregator = currentLineAggregator.concat(wordObj);
      const testLineText = testLineAggregator.map((w) => w.word).join(" ");
      if (
        ctx.measureText(testLineText).width > maxWidth &&
        currentLineAggregator.length > 0
      ) {
        lines.push(currentLineAggregator);
        currentLineAggregator = [wordObj];
      } else {
        currentLineAggregator.push(wordObj);
      }
    }
    if (currentLineAggregator.length > 0) {
      lines.push(currentLineAggregator);
    }
    // If greedy approach still failed (e.g. single word too long), push all as one line.
    if (lines.length === 0 && allWordsWithFormatting.length > 0) {
      lines.push(allWordsWithFormatting);
    }
  }

  // Final cleanup: ensure at most 2 lines if that's a hard requirement.
  // For this problem, we assume 2 lines is the target for problematic splits.
  if (lines.length > 2) {
    console.warn(
      `Splitting resulted in ${lines.length} lines. Forcing to 2 lines by merging subsequent lines.`
    );
    const line1 = lines[0];
    const restCombined = [].concat(...lines.slice(1));
    lines.length = 0;
    if (line1 && line1.length > 0) lines.push(line1);
    if (restCombined.length > 0) lines.push(restCombined);
  }

  // Ensure no empty lines in the result
  return lines.filter((line) => line.length > 0);
}

function drawCaptionWithActiveWord(ctx, caption, activeWordIndex) {
  if (!caption || !caption.words || caption.words.length === 0) return;

  // Set font for measurement and drawing
  ctx.font = `bold ${FONT_SIZE}px Metropolis, Arial, sans-serif`;

  const maxWidth = WIDTH * 0.8; // 80% of canvas width
  const MAX_CHARS_PER_LINE_SINGLE = 15; // MODIFIED: Was 20, now 15 for more aggressive splitting

  // Prepare word data for the splitting function
  const wordsToSplit = caption.words.map((wordData, idx) => ({
    word: wordData.word, // Keep raw word, formatting is handled by splitCaptionIntoLines via formatTextFunc
    originalWordIndex: idx,
    lineBreak: wordData.lineBreak,
    highlighted: wordData.highlighted,
    upperCase: wordData.upperCase,
    // Pass explicit line break hint
  }));

  // Get lines using the new splitting function
  const linesOfWordObjects = splitCaptionIntoLines(
    wordsToSplit,
    ctx,
    maxWidth,
    MAX_CHARS_PER_LINE_SINGLE,
    formatText
  );

  if (linesOfWordObjects.length === 0) {
    // console.warn("Caption resulted in no lines: ", caption.text);
    return; // Nothing to draw
  }

  // Calculate total height of all lines
  const totalHeight = linesOfWordObjects.length * LINE_HEIGHT;

  // Calculate maximum line width for background
  let maxLineWidth = 0;
  for (const line of linesOfWordObjects) {
    const lineText = line
      .map((wo) => (wo.upperCase ? wo.word.toUpperCase() : wo.word))
      .join(" "); // wo.word is already formatted text
    maxLineWidth = Math.max(maxLineWidth, ctx.measureText(lineText).width);
  }

  // Account for shadow and padding
  const shadowMargin =
    Math.max(ACTIVE_SHADOW_OFFSET_X, ACTIVE_SHADOW_OFFSET_Y) * 2 +
    ACTIVE_SHADOW_BLUR;
  const bgWidth = Math.min(
    maxLineWidth + PADDING_X * 2 + shadowMargin,
    WIDTH * 0.9
  );
  const bgHeight = totalHeight + PADDING_Y * 2 + shadowMargin;

  // Draw background rectangle
  const bgX = (WIDTH - bgWidth) / 2;
  const bgY = (HEIGHT - bgHeight) / 2;
  ctx.fillStyle = BG_COLOR;
  drawRoundedRect(ctx, bgX, bgY, bgWidth, bgHeight, BORDER_RADIUS);

  // Calculate proper starting Y position with padding
  // This ensures the text starts after the top padding rather than centered
  let yPosition = bgY + PADDING_Y + FONT_SIZE; // Use font size instead of LINE_HEIGHT/2 to properly align

  for (const line of linesOfWordObjects) {
    // line is an array of {word, originalIndex, ...}
    // Reset text alignment for word-by-word rendering
    ctx.textAlign = "left";

    // Calculate line width for centering
    let lineWidth = 0;
    for (const wordObj of line) {
      lineWidth += ctx.measureText(
        wordObj.upperCase ? wordObj.word.toUpperCase() : wordObj.word
      ).width;
    }
    lineWidth += ctx.measureText(" ").width * (line.length - 1);

    // Start X position for centered line
    let xPosition = bgX + (bgWidth - lineWidth) / 2;

    // Draw each word
    for (const wordObj of line) {
      // wordObj is {word (formatted), originalIndex, ...}
      const isActive = wordObj.originalIndex === activeWordIndex;
      const isUpperCase = wordObj.upperCase;
      const isHighlighted = wordObj.highlighted;
      console.log;
      // Set shadow based on active state
      if (isActive) {
        ctx.shadowColor = SHADOW_COLOR;
        ctx.shadowBlur = ACTIVE_SHADOW_BLUR;
        ctx.shadowOffsetX = ACTIVE_SHADOW_OFFSET_X;
        ctx.shadowOffsetY = ACTIVE_SHADOW_OFFSET_Y;
      } else {
        ctx.shadowColor = SHADOW_COLOR;
        ctx.shadowBlur = INACTIVE_SHADOW_BLUR;
        ctx.shadowOffsetX = INACTIVE_SHADOW_OFFSET_X;
        ctx.shadowOffsetY = INACTIVE_SHADOW_OFFSET_Y;
      }
      // Set text color
      if (isHighlighted && isActive) {
        console.log("selecting highlighted color");
        ctx.fillStyle = ACTIVE_HIGHLIGHT_COLOR;
      } else if (isHighlighted && !isActive) {
        console.log("selecting inactive highlighted color");
        ctx.fillStyle = INACTIVE_HIGHLIGHT_COLOR;
      } else {
        // Set text color
        ctx.fillStyle = isActive ? ACTIVE_COLOR : INACTIVE_COLOR;
      }

      // Draw word
      ctx.fillText(
        isUpperCase ? wordObj.word.toUpperCase() : wordObj.word,
        xPosition,
        yPosition
      );

      // Move X position for next word
      xPosition += ctx.measureText(
        isUpperCase ? wordObj.word.toUpperCase() + " " : wordObj.word + " "
      ).width;
    }

    // Reset shadows
    ctx.shadowColor = "transparent";
    ctx.shadowBlur = 0;
    ctx.shadowOffsetX = 0;
    ctx.shadowOffsetY = 0;

    // Move to next line
    yPosition += LINE_HEIGHT;
  }
}

// Special function to debug file system issues that might happen with the first word
function debugFilesystem(outputDir, word, isFirstWord) {
  if (!isFirstWord) return;

  try {
    console.log("\n=== FILESYSTEM DEBUG ===");
    console.log(`Output directory: ${outputDir}`);
    console.log(`Output dir exists: ${fs.existsSync(outputDir)}`);

    if (fs.existsSync(outputDir)) {
      console.log(
        `Output dir is writable: ${
          fs.accessSync(outputDir, fs.constants.W_OK) === undefined
        }`
      );

      // List the contents of the output directory
      console.log("Files in output directory:");
      const files = fs.readdirSync(outputDir);
      console.log(
        files.slice(0, 10).join(", ") + (files.length > 10 ? "..." : "")
      );

      // Check if we can create a test file
      const testPath = path.join(outputDir, "test_file.txt");
      fs.writeFileSync(testPath, "Test content");
      const testExists = fs.existsSync(testPath);
      console.log(`Test file created successfully: ${testExists}`);
      if (testExists) {
        fs.unlinkSync(testPath);
      }
    }

    console.log("=== END FILESYSTEM DEBUG ===\n");
  } catch (error) {
    console.error(`Filesystem debug error: ${error.message}`);
  }
}

// Start processing
console.log("Generating one PNG per word...");

let lastFinalizedEndFrame = 0; // Tracks the end frame of the previously processed word

async function generateWordCaptions(captions) {
  const canvas = createCanvas(WIDTH, HEIGHT);
  const ctx = canvas.getContext("2d");
  let totalWords = 0;
  if (!generateFirstWordOnly) {
    captions.forEach((caption) => {
      if (caption.words && Array.isArray(caption.words)) {
        totalWords += caption.words.length;
      }
    });
  } else {
    totalWords = 1;
  }
  console.log(
    `Found ${totalWords} words to process${
      generateFirstWordOnly ? " (first word only mode)" : ""
    }`
  );
  let processedWords = 0;

  for (const caption of captions) {
    if (!caption.words || caption.words.length === 0) {
      console.log(`Caption has no word-level timing; skipping`);
      continue;
    }
    if (caption.end < startTime || caption.start > endTime) {
      continue;
    }

    for (let i = 0; i < caption.words.length; i++) {
      const word = caption.words[i];
      const isFirstWordOfAll = processedWords === 0 && generateFirstWordOnly; // Special flag for the absolute first word in generateFirstWordOnly mode
      const isGlobalFirstWord = processedWords === 0; // General first word

      if (generateFirstWordOnly && !isGlobalFirstWord && i > 0) {
        // if first word only, and we already processed it (i.e. on a subsequent word)
        if (isFirstWordOfAll) {
          /* continue if it's not the very first word being generated */
        } else {
          continue;
        }
      }

      if (
        !isGlobalFirstWord &&
        word.end < startTime &&
        !generateFirstWordOnly
      ) {
        // Don't skip if it's the very first word overall.
        continue;
      }
      if (word.start > endTime && !generateFirstWordOnly) {
        continue;
      }

      const MIN_SPAN_FRAMES = 3; // Word must be visible for at least this many frames. (e.g., 3 frames = start, start+1, start+2)

      // 1. Determine current word's START frame
      let uStartFrame = Math.max(1, Math.round(word.start * FPS));
      if (lastFinalizedEndFrame > 0) {
        // Ensure current word starts after the previous one ended, or at its natural start if later.
        uStartFrame = Math.max(uStartFrame, lastFinalizedEndFrame + 1);
      }

      // 2. Determine current word's initial END frame
      let uEndFrame = Math.max(1, Math.round(word.end * FPS)); // Natural end frame from timing
      // Ensure minimum display duration
      uEndFrame = Math.max(uEndFrame, uStartFrame + MIN_SPAN_FRAMES - 1);

      // 3. Gap-filling: If there's a next word IN THE SAME CAPTION, extend current word.
      if (i < caption.words.length - 1) {
        const nextWordInCaption = caption.words[i + 1];
        const nextWordActualStartFrame = Math.max(
          1,
          Math.round(nextWordInCaption.start * FPS)
        );

        // If there's a spoken gap before the next word (next word starts after current word's uEndFrame + 1 frame)
        if (nextWordActualStartFrame > uEndFrame + 1) {
          const prev_uEndFrame_for_log = uEndFrame;
          uEndFrame = nextWordActualStartFrame - 1; // Extend current word to fill the gap
          console.log(
            `Extended word "${word.word}" (frames ${uStartFrame}-${prev_uEndFrame_for_log} to ${uStartFrame}-${uEndFrame}) to fill gap before next word "${nextWordInCaption.word}" (starts ${nextWordActualStartFrame})`
          );
          // Re-ensure min duration for current word after extension
          uEndFrame = Math.max(uEndFrame, uStartFrame + MIN_SPAN_FRAMES - 1);
        }
      }

      // Ensure uEndFrame is not before uStartFrame (can happen if timings are odd or pushed aggressively)
      uEndFrame = Math.max(uEndFrame, uStartFrame);

      processedWords++;
      const percentComplete = ((processedWords / totalWords) * 100).toFixed(1);
      console.log(
        `Rendering ${
          word.highlighted ? "highlighted" : "normal"
        } word ${processedWords}/${totalWords} (${percentComplete}%): "${
          word.word
        }" (frames ${uStartFrame}-${uEndFrame}, time ${word.start.toFixed(
          3
        )}s-${word.end.toFixed(3)}s)${
          isGlobalFirstWord && generateFirstWordOnly
            ? " [FIRST WORD ONLY MODE]"
            : ""
        }`
      );

      if (isGlobalFirstWord && generateFirstWordOnly) {
        debugFilesystem(outputDir, word, true);
      }

      ctx.fillStyle = "rgba(0, 0, 0, 0)";
      ctx.clearRect(0, 0, WIDTH, HEIGHT);
      drawCaptionWithActiveWord(ctx, caption, i);
      const outputFilename = `word_${uStartFrame}_${uEndFrame}.png`;
      const outputPath = path.join(outputDir, outputFilename);
      const buffer = canvas.toBuffer("image/png");
      try {
        fs.writeFileSync(outputPath, buffer);
        if (isGlobalFirstWord && generateFirstWordOnly) {
          console.log(`FILE WRITTEN (First Word Only Mode): ${outputPath}`);
          const fileExists = fs.existsSync(outputPath);
          console.log(`FILE EXISTS CHECK: ${fileExists ? "YES" : "NO"}`);
          console.log(
            `FILE SIZE: ${
              fileExists ? fs.statSync(outputPath).size : "N/A"
            } bytes`
          );
        }
      } catch (error) {
        console.error(
          `Error saving PNG for word "${word.word}": ${error.message}`
        );
      }

      lastFinalizedEndFrame = uEndFrame; // Update for the next word

      if (generateFirstWordOnly && isGlobalFirstWord) {
        console.log(
          "First word generated successfully in first-word-only mode. Exiting."
        );
        return;
      }
    }
  }
}

// Start processing
generateWordCaptions(captionsData.captions)
  .then(() => {
    console.log(
      `\nEach PNG is named with its frame range: word_startFrame_endFrame.png`
    );
    console.log(
      `You can now convert these to MOV files with alpha channels for your timeline.`
    );
  })
  .catch((err) => {
    console.error("Error during generateWordCaptions:", err);
  });