liaoxin-skills/skills/gemini-web-generate/scripts/src/screenshot.js

import path from 'path';
import fs from 'fs';
import { config } from '../config.js';
import { log, emit } from './browser.js';

const sleep = (ms) => new Promise(r => setTimeout(r, ms));

/**
 * Take a screenshot of the current Gemini page.
 * @param {import('puppeteer-core').Page} page
 * @param {object} options
 * @param {boolean} [options.full] - full page screenshot
 * @param {string} [options.outputPath] - custom output path
 * @returns {Promise<{path: string}>}
 */
export async function takeScreenshot(page, { full = false, outputPath, quality = 60 } = {}) {
  fs.mkdirSync(config.screenshotDir, { recursive: true });

  const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
  const filename = `screenshot-${timestamp}.jpg`;
  const filePath = outputPath || path.join(config.screenshotDir, filename);

  await page.screenshot({
    path: filePath,
    fullPage: full,
    type: 'jpeg',
    quality,
  });

  return { path: filePath };
}

/**
 * Count existing download buttons and return their count.
 * @param {import('puppeteer-core').Page} page
 * @returns {Promise<{count: number}>}
 */
export async function snapshotDownloadButtons(page) {
  const buttons = await page.$$('button.generated-image-button');
  const visible = [];
  for (const btn of buttons) {
    if (await btn.isVisible()) {
      visible.push(btn);
    }
  }
  return { count: visible.length };
}

/**
 * Download generated images by clicking download buttons.
 * Waits for all images to finish loading before clicking, then scrolls
 * gently to each button one at a time.
 * @param {import('puppeteer-core').Page} page
 * @param {import('puppeteer-core').CDPSession} cdp
 * @param {object} options
 * @param {number} [options.existingButtonCount] - number of download buttons before this generation
 * @param {number} [options.timeout] - max wait per download in ms (default 120000)
 * @returns {Promise<Array<{path: string}>>}
 */
export async function downloadViaButtons(page, cdp, { existingButtonCount = 0, newestOnly = false, timeout = 120000 } = {}) {
  fs.mkdirSync(config.downloadDir, { recursive: true });

  // Get list of files before download
  const filesBefore = new Set(fs.readdirSync(config.downloadDir));

  // Wait for all images on the page to finish loading and decoding.
  // This is critical: Gemini generates multiple large images, and interacting
  // with the page while they're still decoding can cause tab crashes.
  log('Waiting for all images to finish loading...');
  try {
    await page.evaluate(async () => {
      const images = Array.from(document.images);
      const promises = images
        .filter(img => img.src && !img.complete)
        .map(img => img.decode().catch(() => {}));
      await Promise.all(promises);
    });
    emit({ type: 'progress', step: 'images_ready', message: 'All images loaded' });
  } catch {
    log('Image loading check timed out, proceeding anyway');
  }

  // Find all download buttons — both visible and in the DOM
  const allBtns = await page.$$('button.generated-image-button');

  if (allBtns.length === 0) {
    throw new Error(
      'No generated image download button found. ' +
      'Make sure an image has been generated first.'
    );
  }

  // Determine which buttons to click
  let targetBtns;
  if (newestOnly) {
    // Only click the very last (newest) button
    targetBtns = allBtns.slice(-1);
  } else {
    // Skip the ones that existed before generation
    targetBtns = allBtns.slice(existingButtonCount);
  }

  if (targetBtns.length === 0) {
    log('No new download buttons detected (existing count:', existingButtonCount + ')');
    return [];
  }

  log(`Clicking ${targetBtns.length} download button(s) (total in DOM: ${allBtns.length})`);
  emit({ type: 'progress', step: 'download', buttonCount: targetBtns.length,
    message: `Clicking ${targetBtns.length} download button(s)` });

  // Click each button and wait for the download to complete
  const results = [];
  for (let i = 0; i < targetBtns.length; i++) {
    const btn = targetBtns[i];
    log(`Clicking download button ${i + 1}/${targetBtns.length}...`);

    // Gently scroll the button into view before clicking
    await btn.evaluate(el => el.scrollIntoView({ behavior: 'auto', block: 'center' }));
    await sleep(500);

    await btn.click();

    // Poll for new file in download directory
    const downloadPath = await new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        reject(new Error(`Download timed out after ${timeout / 1000}s`));
      }, timeout);

      const poll = setInterval(() => {
        const filesAfter = fs.readdirSync(config.downloadDir);
        const newFiles = filesAfter.filter(f => !filesBefore.has(f) && !f.endsWith('.crdownload'));
        if (newFiles.length > 0) {
          clearInterval(poll);
          clearTimeout(timer);
          // Wait a bit for file to finish writing
          setTimeout(() => {
            // Update filesBefore so next download is tracked separately
            for (const f of newFiles) {
              filesBefore.add(f);
            }
            resolve(path.join(config.downloadDir, newFiles[0]));
          }, 1000);
        }
      }, 500);
    });

    results.push({ path: downloadPath });
    log(`Downloaded: ${downloadPath}`);

    // Small delay between downloads
    await sleep(500);
  }

  return results;
}