Mycontent/lib/documentation/screenshot.ts

import type { DocumentationPage } from "@/lib/documentation/catalog"

export type ScreenshotResult = {
  page: DocumentationPage
  screenshotPng?: Buffer
  error?: string
}

type CaptureOptions = {
  baseUrl: string
  timeoutMs: number
  viewport: { width: number; height: number }
}

export async function captureScreenshots(
  pages: DocumentationPage[],
  options: CaptureOptions,
): Promise<ScreenshotResult[]> {
  const { chromium } = await import("playwright")
  const browser = await chromium.launch({
    headless: true,
    args: ["--no-sandbox", "--disable-setuid-sandbox"],
  })

  try {
    const results: ScreenshotResult[] = []

    for (const pageInfo of pages) {
      const page = await browser.newPage({ viewport: options.viewport })
      try {
        const captureUrl = new URL("/documentation/capture", options.baseUrl)
        captureUrl.searchParams.set("path", pageInfo.path)

        console.log(`[v0] Capturing: ${pageInfo.path}`)

        await page.goto(captureUrl.toString(), {
          waitUntil: "networkidle",
          timeout: options.timeoutMs,
        })

        const iframeHandle = await page.waitForSelector('iframe[data-doc-iframe="true"]', {
          timeout: options.timeoutMs,
        })

        const frame = await iframeHandle.contentFrame()
        if (!frame) {
          throw new Error("无法获取iframe内容")
        }

        await frame.waitForLoadState("domcontentloaded", { timeout: options.timeoutMs })

        // Allow network to settle
        await frame.waitForLoadState("networkidle", { timeout: options.timeoutMs }).catch(() => {
          console.log(`[v0] Network idle timeout for ${pageInfo.path}, continuing...`)
        })

        if (pageInfo.waitForSelector) {
          await frame
            .waitForSelector(pageInfo.waitForSelector, {
              timeout: options.timeoutMs,
            })
            .catch(() => {
              console.log(`[v0] Selector timeout for ${pageInfo.path}, continuing...`)
            })
        }

        await page.waitForTimeout(500)

        const screenshot = await iframeHandle.screenshot({
          type: "png",
          animations: "disabled",
        })

        results.push({
          page: pageInfo,
          screenshotPng: Buffer.from(screenshot),
        })

        console.log(`[v0] Success: ${pageInfo.path}`)
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error)
        console.log(`[v0] Error capturing ${pageInfo.path}: ${message}`)
        results.push({ page: pageInfo, error: message })
      } finally {
        await page.close().catch(() => undefined)
      }
    }

    return results
  } finally {
    await browser.close().catch(() => undefined)
  }
}