Build a Screenshot API

// src/screenshot/service.ts — Screenshot API with browser pooling, caching, and queue
import puppeteer, { Browser, Page } from "puppeteer";
import { Redis } from "ioredis";
import { createHash } from "node:crypto";
import sharp from "sharp";

const redis = new Redis(process.env.REDIS_URL!);

let browser: Browser;
const MAX_PAGES = 5;           // concurrent tabs
const PAGE_TIMEOUT = 15000;
let activeTabs = 0;

interface ScreenshotRequest {
  url: string;
  viewport?: { width: number; height: number };
  fullPage?: boolean;
  selector?: string;           // capture specific element
  format?: "png" | "jpeg" | "webp";
  quality?: number;            // 1-100 for jpeg/webp
  scale?: number;              // device scale factor (1-3)
  delay?: number;              // ms to wait after load
  blockAds?: boolean;
  darkMode?: boolean;
  hideSelectors?: string[];    // CSS selectors to hide
  clip?: { x: number; y: number; width: number; height: number };
  maxWidth?: number;           // resize output
}

interface ScreenshotResult {
  buffer: Buffer;
  width: number;
  height: number;
  format: string;
  size: number;
  cached: boolean;
  url: string;
}

// Initialize browser
export async function initBrowser(): Promise<void> {
  browser = await puppeteer.launch({
    headless: true,
    args: [
      "--no-sandbox", "--disable-setuid-sandbox",
      "--disable-dev-shm-usage", "--disable-gpu",
      "--single-process", "--no-zygote",
      "--disable-web-security",
    ],
  });
}

// Take screenshot
export async function takeScreenshot(req: ScreenshotRequest): Promise<ScreenshotResult> {
  // Check cache
  const cacheKey = `ss:${createHash("md5").update(JSON.stringify(req)).digest("hex")}`;
  const cached = await redis.getBuffer(cacheKey);
  if (cached) {
    const meta = await redis.get(`${cacheKey}:meta`);
    const { width, height, format } = JSON.parse(meta || "{}");
    return { buffer: cached, width, height, format, size: cached.length, cached: true, url: req.url };
  }

  // Wait for available tab
  while (activeTabs >= MAX_PAGES) {
    await new Promise((r) => setTimeout(r, 100));
  }
  activeTabs++;

  let page: Page | null = null;
  try {
    page = await browser.newPage();

    const viewport = req.viewport || { width: 1280, height: 720 };
    await page.setViewport({
      width: viewport.width,
      height: viewport.height,
      deviceScaleFactor: req.scale || 1,
    });

    // Dark mode
    if (req.darkMode) {
      await page.emulateMediaFeatures([{ name: "prefers-color-scheme", value: "dark" }]);
    }

    // Block ads/trackers
    if (req.blockAds) {
      await page.setRequestInterception(true);
      page.on("request", (request) => {
        const blocked = ["google-analytics", "facebook", "doubleclick", "ads"];
        if (blocked.some((b) => request.url().includes(b))) {
          request.abort();
        } else {
          request.continue();
        }
      });
    }

    await page.goto(req.url, { waitUntil: "networkidle2", timeout: PAGE_TIMEOUT });

    // Additional delay
    if (req.delay) await new Promise((r) => setTimeout(r, req.delay));

    // Hide elements
    if (req.hideSelectors?.length) {
      for (const sel of req.hideSelectors) {
        await page.evaluate((s) => {
          document.querySelectorAll(s).forEach((el) => {
            (el as HTMLElement).style.display = "none";
          });
        }, sel);
      }
    }

    // Capture
    const format = req.format || "png";
    let screenshotOptions: any = {
      type: format === "webp" ? "png" : format, // puppeteer doesn't support webp directly
      fullPage: req.fullPage || false,
      quality: format === "jpeg" ? (req.quality || 80) : undefined,
    };

    if (req.selector) {
      const element = await page.$(req.selector);
      if (!element) throw new Error(`Selector not found: ${req.selector}`);
      screenshotOptions.clip = await element.boundingBox();
    } else if (req.clip) {
      screenshotOptions.clip = req.clip;
    }

    let buffer = await page.screenshot(screenshotOptions) as Buffer;

    // Post-process with sharp
    let image = sharp(buffer);

    if (format === "webp") {
      image = image.webp({ quality: req.quality || 80 });
    }

    if (req.maxWidth) {
      const metadata = await sharp(buffer).metadata();
      if (metadata.width && metadata.width > req.maxWidth) {
        image = image.resize(req.maxWidth);
      }
    }

    buffer = await image.toBuffer();
    const metadata = await sharp(buffer).metadata();

    const result: ScreenshotResult = {
      buffer,
      width: metadata.width || viewport.width,
      height: metadata.height || viewport.height,
      format,
      size: buffer.length,
      cached: false,
      url: req.url,
    };

    // Cache for 1 hour
    await redis.setex(cacheKey, 3600, buffer);
    await redis.setex(`${cacheKey}:meta`, 3600, JSON.stringify({
      width: result.width, height: result.height, format,
    }));

    return result;
  } finally {
    if (page) await page.close().catch(() => {});
    activeTabs--;
  }
}

// Batch screenshots
export async function batchScreenshots(
  requests: ScreenshotRequest[]
): Promise<ScreenshotResult[]> {
  // Process in parallel up to MAX_PAGES
  const results: ScreenshotResult[] = [];
  const chunks = [];
  for (let i = 0; i < requests.length; i += MAX_PAGES) {
    chunks.push(requests.slice(i, i + MAX_PAGES));
  }

  for (const chunk of chunks) {
    const chunkResults = await Promise.all(chunk.map((req) => takeScreenshot(req)));
    results.push(...chunkResults);
  }

  return results;
}