bytefun/src/screenshotCapture.ts

一个打通了原型设计、UI设计与代码转换、跨平台原生代码开发等的平台

import * as vscode from 'vscode';
import * as fs from 'fs';
import * as path from 'path';
import puppeteer, { Browser, Page, HTTPRequest } from 'puppeteer';

/**
 * HTML页面截图捕获器
 * 使用Puppeteer实现高质量的HTML页面截图功能
 */
export class ScreenshotCapture {
    private static instance: ScreenshotCapture;
    private browser: Browser | null = null;
    private isInitializing: boolean = false;
    private initPromise: Promise<Browser> | null = null;

    private constructor() { }

    /**
     * 获取单例实例
     */
    public static getInstance(): ScreenshotCapture {
        if (!ScreenshotCapture.instance) {
            ScreenshotCapture.instance = new ScreenshotCapture();
        }
        return ScreenshotCapture.instance;
    }

    /**
     * 初始化Puppeteer浏览器实例
     */
    private async initializeBrowser(): Promise<Browser> {
        if (this.browser && this.browser.isConnected()) {
            return this.browser;
        }

        // 如果正在初始化，等待初始化完成
        if (this.isInitializing && this.initPromise) {
            return this.initPromise;
        }

        this.isInitializing = true;


        this.initPromise = puppeteer.launch({
            headless: true,
            defaultViewport: null,
            args: [
                '--no-sandbox',
                '--disable-setuid-sandbox',
                '--disable-dev-shm-usage',
                '--disable-gpu',
                '--no-first-run',
                '--disable-background-timer-throttling',
                '--disable-backgrounding-occluded-windows',
                '--disable-renderer-backgrounding',
                '--disable-features=TranslateUI',
                '--disable-ipc-flooding-protection',
                '--window-size=1920,1080'
            ],
            ignoreDefaultArgs: ['--disable-extensions'],
            timeout: 30000
        });

        try {
            this.browser = await this.initPromise;


            // 监听浏览器断开连接
            this.browser.on('disconnected', () => {

                this.browser = null;
                this.isInitializing = false;
                this.initPromise = null;
            });

            return this.browser;
        } catch (error) {
            console.error('❌ [ScreenshotCapture] Puppeteer浏览器初始化失败:', error);
            this.isInitializing = false;
            this.initPromise = null;
            throw error;
        } finally {
            this.isInitializing = false;
        }
    }

    /**
     * 捕获HTML页面截图
     * @param htmlFilePath HTML文件的完整路径
     * @param workspaceRoot 工作区根目录
     * @param options 截图选项
     * @returns 截图文件路径，失败返回null
     */
    public async captureHtmlPage(
        htmlFilePath: string,
        workspaceRoot: string,
        options: ScreenshotOptions = {}
    ): Promise<ScreenshotResult> {
        const startTime = Date.now();


        try {
            // 验证HTML文件是否存在
            if (!fs.existsSync(htmlFilePath)) {
                throw new Error(`HTML文件不存在: ${htmlFilePath}`);
            }

            // 确保截图目录存在
            const screenshotDir = path.join(workspaceRoot, 'screenshot');
            if (!fs.existsSync(screenshotDir)) {
                fs.mkdirSync(screenshotDir, { recursive: true });

            }

            // 生成截图文件名
            const fileName = this.generateScreenshotFileName(htmlFilePath, workspaceRoot, options.suffix);
            const screenshotPath = path.join(screenshotDir, fileName);

            // 初始化浏览器
            const browser = await this.initializeBrowser();

            // 创建新页面
            const page = await browser.newPage();

            try {
                // 配置页面
                await this.configurePage(page, options);

                // 加载HTML文件
                const fileUrl = `file://${htmlFilePath}`;


                await page.goto(fileUrl, {
                    waitUntil: 'networkidle0',
                    timeout: options.timeout || 30000
                });

                // 等待页面完全渲染
                await this.waitForPageReady(page, options);

                // 判断是否需要截取特定元素
                if (options.elementSelector) {

                    await this.captureElementScreenshot(page, screenshotPath, options);
                } else {
                    // 执行全页面截图
                    await page.screenshot({
                        path: screenshotPath,
                        fullPage: options.fullPage !== false,
                        type: 'png',
                        omitBackground: options.transparent || false,
                        clip: options.clip
                    });
                }

                const duration = Date.now() - startTime;


                return {
                    success: true,
                    screenshotPath,
                    relativePath: path.relative(workspaceRoot, screenshotPath),
                    duration
                };

            } finally {
                // 确保页面被关闭
                try {
                    await page.close();
                } catch (error) {
                    console.warn('⚠️ [ScreenshotCapture] 关闭页面时出错:', error);
                }
            }

        } catch (error) {
            const duration = Date.now() - startTime;
            console.error('❌ [ScreenshotCapture] 截图失败:', error);

            return {
                success: false,
                error: error instanceof Error ? error.message : String(error),
                duration
            };
        }
    }

    /**
     * 捕获特定元素的截图
     */
    private async captureElementScreenshot(page: Page, screenshotPath: string, options: ScreenshotOptions): Promise<void> {
        try {


            // 等待元素出现
            await page.waitForSelector(options.elementSelector!, {
                timeout: options.selectorTimeout || 10000
            });

            // 获取元素
            const element = await page.$(options.elementSelector!);
            if (!element) {
                throw new Error(`未找到元素: ${options.elementSelector}`);
            }



            // 直接对元素进行截图
            await element.screenshot({
                path: screenshotPath,
                type: 'png',
                omitBackground: options.transparent || false
            });



        } catch (error) {
            console.error(`❌ [ScreenshotCapture] 元素截图失败:`, error);

            // 如果元素截图失败，回退到全页面截图

            await page.screenshot({
                path: screenshotPath,
                fullPage: options.fullPage !== false,
                type: 'png',
                omitBackground: options.transparent || false,
                clip: options.clip
            });
        }
    }

    /**
     * 配置页面参数
     */
    private async configurePage(page: Page, options: ScreenshotOptions): Promise<void> {
        // 设置视窗大小
        await page.setViewport({
            width: options.width || 1920,
            height: options.height || 1080,
            deviceScaleFactor: options.deviceScaleFactor || 2
        });

        // 设置用户代理
        await page.setUserAgent(
            'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
        );

        // 拦截不必要的资源请求以提高性能
        if (options.blockResources) {
            await page.setRequestInterception(true);
            page.on('request', (request: HTTPRequest) => {
                const resourceType = request.resourceType();
                const url = request.url();

                // 只拦截外部的字体和图片，保留本地文件
                if (['font', 'image'].includes(resourceType) && !url.startsWith('file://') && !url.startsWith('data:')) {

                    request.abort();
                } else {
                    request.continue();
                }
            });
        }

        // 禁用JavaScript（如果需要）
        if (options.disableJavaScript) {
            await page.setJavaScriptEnabled(false);
        }
    }

    /**
     * 等待页面完全加载
     */
    private async waitForPageReady(page: Page, options: ScreenshotOptions): Promise<void> {
        // 基础等待时间
        const baseWaitTime = options.waitTime || 1000;
        await page.waitForTimeout(baseWaitTime);

        // 等待特定元素（如果指定）
        if (options.waitForSelector) {
            try {
                await page.waitForSelector(options.waitForSelector, {
                    timeout: options.selectorTimeout || 5000
                });
            } catch (error) {
                console.warn(`⚠️ [ScreenshotCapture] 等待选择器超时: ${options.waitForSelector}`);
            }
        }

        // 如果指定了元素选择器，也要等待该元素
        if (options.elementSelector && options.elementSelector !== options.waitForSelector) {
            try {
                await page.waitForSelector(options.elementSelector, {
                    timeout: options.selectorTimeout || 5000
                });

            } catch (error) {
                console.warn(`⚠️ [ScreenshotCapture] 等待目标元素超时: ${options.elementSelector}`);
            }
        }

        // 等待字体加载完成
        try {
            await page.evaluate(() => {
                return (document as any).fonts?.ready || Promise.resolve();
            });
        } catch (error) {
            console.warn('⚠️ [ScreenshotCapture] 等待字体加载失败:', error);
        }

        // 等待所有图片加载完成
        try {
            await page.evaluate(() => {
                return Promise.all(
                    Array.from(document.images).map((img: HTMLImageElement) => {
                        if (img.complete) {
                            return Promise.resolve();
                        }
                        return new Promise((resolve, reject) => {
                            img.addEventListener('load', resolve);
                            img.addEventListener('error', resolve); // 即使加载失败也继续
                            // 5秒超时
                            setTimeout(resolve, 5000);
                        });
                    })
                );
            });

        } catch (error) {
            console.warn('⚠️ [ScreenshotCapture] 等待图片加载失败:', error);
        }

        // 额外等待确保渲染完成
        await page.waitForTimeout(1000);
    }

    /**
     * 生成截图文件名
     */
    private generateScreenshotFileName(htmlFilePath: string, workspaceRoot: string, suffix?: string): string {
        const baseName = path.basename(htmlFilePath, '.html');

        // 从UI设计进度.json文件中查找序号
        const sequenceNumber = this.findSequenceNumberFromUIProgress(baseName, workspaceRoot);

        // 如果找到序号，使用新格式：序号-html文件名.png
        if (sequenceNumber !== null) {
            const suffixPart = suffix ? `_${suffix}` : '';
            return `${sequenceNumber}-${baseName}${suffixPart}.png`;
        }

        // 如果没找到序号，使用原来的格式作为fallback
        const timestamp = new Date().toISOString()
            .replace(/:/g, '-')
            .replace(/\..+/, '')
            .replace('T', '_');

        const suffixPart = suffix ? `_${suffix}` : '';
        return `${baseName}${suffixPart}_${timestamp}.png`;
    }

    /**
     * 从UI设计进度.json文件中查找对应的序号
     */
    private findSequenceNumberFromUIProgress(htmlFileName: string, workspaceRoot: string): number | null {
        try {


            // 构建UI设计进度.json文件路径
            const progressFilePath = path.join(workspaceRoot, 'doc', 'UI设计进度.json');

            // 检查文件是否存在
            if (!fs.existsSync(progressFilePath)) {

                return null;
            }

            // 读取文件内容
            const content = fs.readFileSync(progressFilePath, 'utf8');


            // 解析JSON格式
            const progressData = JSON.parse(content);


            // 将html文件名小写进行匹配（如adPage -> adpage）
            const capitalizedPageName = htmlFileName.toLowerCase();


            // 遍历所有模块和页面查找匹配
            let pageIndex = 1;

            if (progressData.modules && progressData.modules.length > 0) {
                for (const module of progressData.modules) {
                    if (module.pages && module.pages.length > 0) {
                        for (const page of module.pages) {
                            const pageEnName = page.pageNameEN || page.pageNameCN;
                            const pageNameLower = pageEnName.toLowerCase();



                            // 检查页面名称是否匹配
                            if (pageNameLower === capitalizedPageName) {

                                return pageIndex;
                            }

                            pageIndex++;
                        }
                    }
                }
            }


            return null;

        } catch (error) {
            console.error('❌ [ScreenshotCapture] 解析UI设计进度.json文件失败:', error);
            return null;
        }
    }

    /**
     * 批量截图
     */
    public async captureMultiplePages(
        htmlFiles: string[],
        workspaceRoot: string,
        options: ScreenshotOptions = {}
    ): Promise<ScreenshotResult[]> {


        const results: ScreenshotResult[] = [];

        for (let i = 0; i < htmlFiles.length; i++) {
            const htmlFile = htmlFiles[i];


            const result = await this.captureHtmlPage(htmlFile, workspaceRoot, {
                ...options,
                suffix: `batch_${i + 1}`
            });

            results.push(result);

            // 批量截图时添加短暂延迟，避免资源占用过高
            if (i < htmlFiles.length - 1) {
                await new Promise(resolve => setTimeout(resolve, 100));
            }
        }

        const successCount = results.filter(r => r.success).length;


        return results;
    }

    /**
     * 获取浏览器状态
     */
    public getBrowserStatus(): BrowserStatus {
        return {
            isConnected: this.browser?.isConnected() || false,
            isInitializing: this.isInitializing,
            version: null // 移除版本信息以避免异步问题
        };
    }

    /**
     * 手动清理浏览器实例
     */
    public async cleanup(): Promise<void> {
        if (this.browser) {
            try {

                await this.browser.close();
                this.browser = null;
                this.isInitializing = false;
                this.initPromise = null;

            } catch (error) {
                console.error('❌ [ScreenshotCapture] 清理浏览器实例失败:', error);
            }
        }
    }

    /**
     * 静态清理方法（用于插件卸载）
     */
    public static async dispose(): Promise<void> {
        if (ScreenshotCapture.instance) {
            await ScreenshotCapture.instance.cleanup();
            ScreenshotCapture.instance = null as any;
        }
    }
}

/**
 * 截图选项接口
 */
export interface ScreenshotOptions {
    /** 视窗宽度，默认1920 */
    width?: number;
    /** 视窗高度，默认1080 */
    height?: number;
    /** 设备像素比，默认2 */
    deviceScaleFactor?: number;
    /** 是否截取完整页面，默认true */
    fullPage?: boolean;
    /** 是否透明背景，默认false */
    transparent?: boolean;
    /** 截图区域 */
    clip?: { x: number; y: number; width: number; height: number };
    /** 页面加载超时时间，默认30000ms */
    timeout?: number;
    /** 等待渲染完成的时间，默认1000ms */
    waitTime?: number;
    /** 等待特定选择器 */
    waitForSelector?: string;
    /** 选择器等待超时时间，默认5000ms */
    selectorTimeout?: number;
    /** 文件名后缀 */
    suffix?: string;
    /** 是否禁用JavaScript */
    disableJavaScript?: boolean;
    /** 是否拦截外部资源 */
    blockResources?: boolean;
    /** 特定元素的选择器，用于截取该元素的截图 */
    elementSelector?: string;
}

/**
 * 截图结果接口
 */
export interface ScreenshotResult {
    /** 是否成功 */
    success: boolean;
    /** 截图文件绝对路径 */
    screenshotPath?: string;
    /** 截图文件相对路径 */
    relativePath?: string;
    /** 错误信息 */
    error?: string;
    /** 执行耗时（毫秒） */
    duration: number;
}

/**
 * 浏览器状态接口
 */
export interface BrowserStatus {
    /** 是否已连接 */
    isConnected: boolean;
    /** 是否正在初始化 */
    isInitializing: boolean;
    /** 浏览器版本 */
    version: string | null;
}

/**
 * 便捷的截图函数
 */
export async function captureHtmlScreenshot(
    htmlFilePath: string,
    workspaceRoot: string,
    options?: ScreenshotOptions
): Promise<ScreenshotResult> {
    const capture = ScreenshotCapture.getInstance();
    return capture.captureHtmlPage(htmlFilePath, workspaceRoot, options);
}