Race-Nation-Shopify-Backend/src/business-logic/kyt-pipeline/03_watermark_downloaded_images.js

const fs = require("node:fs/promises");
const path = require("node:path");
const sharp = require("sharp");

const DEFAULT_IMAGES_DIR = path.join("data", "02_downloaded_product_images");
const DEFAULT_WATERMARK_PATH = path.join("data", "watermark.png");
const DEFAULT_STATE_PATH = path.join("data", "03_watermark_state.json");
const WATERMARK_ENGINE_VERSION = 1;

const IMAGE_EXTENSIONS = new Set([
  ".jpg",
  ".jpeg",
  ".png",
  ".webp",
  ".tif",
  ".tiff",
  ".avif"
]);

async function getAllImageFilesRecursively(rootDir) {
  const files = [];

  async function walk(currentDir) {
    const entries = await fs.readdir(currentDir, { withFileTypes: true });

    for (const entry of entries) {
      const fullPath = path.join(currentDir, entry.name);

      if (entry.isDirectory()) {
        await walk(fullPath);
        continue;
      }

      if (!entry.isFile()) {
        continue;
      }

      const ext = path.extname(entry.name).toLowerCase();
      if (IMAGE_EXTENSIONS.has(ext)) {
        files.push(fullPath);
      }
    }
  }

  await walk(rootDir);
  return files;
}

async function readStateFile(statePath) {
  try {
    const raw = await fs.readFile(statePath, "utf8");
    const parsed = JSON.parse(raw);
    if (!parsed || typeof parsed !== "object") {
      return {};
    }
    return parsed;
  } catch {
    return {};
  }
}

async function getFileFingerprint(filePath) {
  const stat = await fs.stat(filePath);
  return `${stat.size}:${Math.trunc(stat.mtimeMs)}`;
}

async function createTiledWatermark(imageBuffer, watermarkPath) {
  const image = sharp(imageBuffer);
  const imageMeta = await image.metadata();
  const width = imageMeta.width || 0;
  const height = imageMeta.height || 0;

  if (width === 0 || height === 0) {
    throw new Error("Image has invalid dimensions");
  }

  const scaledWatermark = await sharp(watermarkPath)
    .resize({ width: Math.floor(Math.min(width, height) / 3) })
    .toBuffer();

  const wmMeta = await sharp(scaledWatermark).metadata();
  const wmWidth = wmMeta.width || 1;
  const wmHeight = wmMeta.height || 1;

  const positions = [];
  for (let y = 0; y < height; y += wmHeight) {
    for (let x = 0; x < width; x += wmWidth) {
      positions.push({ input: scaledWatermark, left: x, top: y });
    }
  }

  return sharp(imageBuffer).composite(positions).toBuffer();
}

async function watermarkImageInPlace(imagePath, watermarkPath) {
  const ext = path.extname(imagePath).toLowerCase();
  let image = sharp(imagePath, { animated: false });

  let meta = await image.metadata();
  let width = meta.width || 0;
  let height = meta.height || 0;
  const format = (meta.format || "").toLowerCase();

  if (width === 0 || height === 0) {
    throw new Error("Image has invalid metadata");
  }

  // Replicates the old logic: minimum size 500x500 before watermarking.
  if (width < 500 || height < 500) {
    const newWidth = width < 500 ? 500 : width;
    const newHeight = height < 500 ? 500 : height;
    image = image.resize(newWidth, newHeight);
    width = newWidth;
    height = newHeight;
  }

  let baseBuffer;

  if (format === "png" || ext === ".png") {
    const background = sharp({
      create: {
        width,
        height,
        channels: 4,
        background: { r: 255, g: 255, b: 255, alpha: 1 }
      }
    });

    baseBuffer = await background
      .composite([{ input: await image.toBuffer(), gravity: "center" }])
      .png()
      .toBuffer();
  } else {
    baseBuffer = await image.jpeg().toBuffer();
  }

  const watermarkedBuffer = await createTiledWatermark(baseBuffer, watermarkPath);

  const writer = sharp(watermarkedBuffer);
  if (ext === ".png") {
    await writer.png().toFile(imagePath);
  } else if (ext === ".webp") {
    await writer.webp().toFile(imagePath);
  } else if (ext === ".avif") {
    await writer.avif().toFile(imagePath);
  } else if (ext === ".tif" || ext === ".tiff") {
    await writer.tiff().toFile(imagePath);
  } else {
    await writer.jpeg().toFile(imagePath);
  }
}

async function applyWatermarkToDownloadedImages(options = {}) {
  const {
    imagesDir = DEFAULT_IMAGES_DIR,
    watermarkPath = DEFAULT_WATERMARK_PATH,
    statePath = DEFAULT_STATE_PATH
  } = options;

  const absImagesDir = path.resolve(process.cwd(), imagesDir);
  const absWatermarkPath = path.resolve(process.cwd(), watermarkPath);
  const absStatePath = path.resolve(process.cwd(), statePath);

  await fs.access(absImagesDir);
  await fs.access(absWatermarkPath);

  await fs.mkdir(path.dirname(absStatePath), { recursive: true });

  const watermarkFingerprint = await getFileFingerprint(absWatermarkPath);
  const state = await readStateFile(absStatePath);
  const stateVersion = Number(state.engineVersion || 0);
  const stateWatermarkFingerprint = String(state.watermarkFingerprint || "");
  const previousFiles = state.files && typeof state.files === "object" ? state.files : {};
  const stateFiles =
    stateVersion === WATERMARK_ENGINE_VERSION && stateWatermarkFingerprint === watermarkFingerprint
      ? { ...previousFiles }
      : {};

  const imageFiles = await getAllImageFilesRecursively(absImagesDir);

  let processed = 0;
  let skipped = 0;
  let failed = 0;

  for (let i = 0; i < imageFiles.length; i += 1) {
    const imagePath = imageFiles[i];
    const relativePath = path.relative(absImagesDir, imagePath);

    try {
      const beforeFingerprint = await getFileFingerprint(imagePath);
      if (stateFiles[relativePath] === beforeFingerprint) {
        skipped += 1;
      } else {
      await watermarkImageInPlace(imagePath, absWatermarkPath);
        const afterFingerprint = await getFileFingerprint(imagePath);
        stateFiles[relativePath] = afterFingerprint;
      processed += 1;
      }
    } catch (error) {
      failed += 1;
      console.log(`[WATERMARK-FAIL] ${imagePath} -> ${error.message}`);
    }

    if ((i + 1) % 50 === 0 || i === imageFiles.length - 1) {
      console.log(`[WATERMARK] ${i + 1}/${imageFiles.length} processed`);
    }
  }

  // Keep state only for files that still exist.
  const currentSet = new Set(imageFiles.map((x) => path.relative(absImagesDir, x)));
  for (const rel of Object.keys(stateFiles)) {
    if (!currentSet.has(rel)) {
      delete stateFiles[rel];
    }
  }

  const finalState = {
    engineVersion: WATERMARK_ENGINE_VERSION,
    watermarkFingerprint,
    updatedAt: new Date().toISOString(),
    files: stateFiles
  };
  await fs.writeFile(absStatePath, JSON.stringify(finalState, null, 2), "utf8");

  return {
    imagesDir: absImagesDir,
    watermarkPath: absWatermarkPath,
    statePath: absStatePath,
    totalImagesFound: imageFiles.length,
    processed,
    skipped,
    failed
  };
}

module.exports = {
  applyWatermarkToDownloadedImages
};

if (require.main === module) {
  applyWatermarkToDownloadedImages()
    .then((summary) => {
      console.log("\nWatermark summary:");
      console.log(JSON.stringify(summary, null, 2));
    })
    .catch((error) => {
      console.error("Watermark run failed:", error.message);
      process.exitCode = 1;
    });
}