package utils import ( "bytes" "cellbot/internal/protocol" "context" "encoding/base64" "fmt" "html/template" "os" "path/filepath" "time" "github.com/chromedp/chromedp" "go.uber.org/zap" ) // ScreenshotOptions 截图选项 type ScreenshotOptions struct { Width int // 视口宽度(像素) Height int // 视口高度(像素) Timeout time.Duration // 超时时间 WaitTime time.Duration // 等待时间(页面加载后等待) FullPage bool // 是否截取整个页面 Quality int // 图片质量(0-100,仅PNG格式) Format string // 图片格式:png, jpeg Logger *zap.Logger // 日志记录器 } // DefaultScreenshotOptions 默认截图选项 func DefaultScreenshotOptions() *ScreenshotOptions { return &ScreenshotOptions{ Width: 1920, Height: 1080, Timeout: 30 * time.Second, WaitTime: 1 * time.Second, FullPage: false, Quality: 90, Format: "png", Logger: zap.NewNop(), } } // ScreenshotURL 对指定URL进行截图并返回base64编码 func ScreenshotURL(ctx context.Context, url string, opts *ScreenshotOptions) (string, error) { if opts == nil { opts = DefaultScreenshotOptions() } // 创建上下文,添加优化选项 allocCtx, cancel := chromedp.NewExecAllocator(ctx, chromedp.NoSandbox, chromedp.NoFirstRun, chromedp.NoDefaultBrowserCheck, chromedp.Headless, chromedp.DisableGPU, ) defer cancel() ctx, cancel = chromedp.NewContext(allocCtx, chromedp.WithLogf(func(format string, v ...interface{}) { if opts.Logger != nil { opts.Logger.Debug(fmt.Sprintf(format, v...)) } })) defer cancel() // 设置超时 ctx, cancel = context.WithTimeout(ctx, opts.Timeout) defer cancel() var buf []byte // 执行截图任务 var err error if opts.FullPage { err = chromedp.Run(ctx, chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height)), // 设置视口大小 chromedp.Navigate(url), chromedp.WaitReady("body", chromedp.ByQuery), // 使用 WaitReady 等待页面完全加载 chromedp.Sleep(opts.WaitTime), chromedp.FullScreenshot(&buf, opts.Quality), ) } else { err = chromedp.Run(ctx, chromedp.Navigate(url), chromedp.WaitReady("body", chromedp.ByQuery), // 使用 WaitReady 等待页面完全加载 chromedp.Sleep(opts.WaitTime), chromedp.CaptureScreenshot(&buf), ) } if err != nil { return "", fmt.Errorf("failed to capture screenshot: %w", err) } // 转换为base64 base64Str := base64.StdEncoding.EncodeToString(buf) return base64Str, nil } // ScreenshotHTML 对HTML内容进行截图并返回base64编码 func ScreenshotHTML(ctx context.Context, htmlContent string, opts *ScreenshotOptions) (string, error) { if opts == nil { opts = DefaultScreenshotOptions() } if opts.Logger != nil { opts.Logger.Info("Starting HTML screenshot", zap.Int("html_length", len(htmlContent)), zap.Int("width", opts.Width), zap.Int("height", opts.Height), zap.Duration("timeout", opts.Timeout)) } // 创建上下文 allocCtx, cancel := chromedp.NewExecAllocator(ctx, chromedp.NoSandbox, chromedp.NoFirstRun, chromedp.NoDefaultBrowserCheck, chromedp.Headless, chromedp.DisableGPU, ) defer cancel() if opts.Logger != nil { opts.Logger.Debug("Chrome allocator created") } ctx, cancel = chromedp.NewContext(allocCtx, chromedp.WithLogf(func(format string, v ...interface{}) { if opts.Logger != nil { opts.Logger.Debug(fmt.Sprintf(format, v...)) } })) defer cancel() // 设置超时 ctx, cancel = context.WithTimeout(ctx, opts.Timeout) defer cancel() var buf []byte // 使用 base64 编码的 data URL,避免 URL 编码导致的 + 号问题 htmlBytes := []byte(htmlContent) htmlBase64 := base64.StdEncoding.EncodeToString(htmlBytes) dataURL := fmt.Sprintf("data:text/html;charset=utf-8;base64,%s", htmlBase64) if opts.Logger != nil { opts.Logger.Debug("Navigating to base64 data URL", zap.Int("html_length", len(htmlContent)), zap.Int("base64_length", len(htmlBase64))) } // 执行截图任务 var err error if opts.FullPage { if opts.Logger != nil { opts.Logger.Debug("Taking full page screenshot") } err = chromedp.Run(ctx, chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height)), // 设置视口大小 chromedp.Navigate(dataURL), chromedp.WaitReady("body", chromedp.ByQuery), chromedp.Sleep(opts.WaitTime), chromedp.FullScreenshot(&buf, opts.Quality), ) } else { if opts.Logger != nil { opts.Logger.Debug("Taking viewport screenshot") } err = chromedp.Run(ctx, chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height)), // 设置视口大小 chromedp.Navigate(dataURL), chromedp.WaitReady("body", chromedp.ByQuery), chromedp.Sleep(opts.WaitTime), chromedp.CaptureScreenshot(&buf), ) } if err != nil { if opts.Logger != nil { opts.Logger.Error("Failed to capture screenshot", zap.Error(err)) } return "", fmt.Errorf("failed to capture screenshot: %w", err) } if opts.Logger != nil { opts.Logger.Info("Screenshot captured successfully", zap.Int("image_size", len(buf))) } // 转换为base64 base64Str := base64.StdEncoding.EncodeToString(buf) if opts.Logger != nil { opts.Logger.Info("Screenshot converted to base64", zap.Int("base64_length", len(base64Str))) } return base64Str, nil } // ScreenshotElement 对页面中的特定元素进行截图 func ScreenshotElement(ctx context.Context, url string, selector string, opts *ScreenshotOptions) (string, error) { if opts == nil { opts = DefaultScreenshotOptions() } // 创建上下文 allocCtx, cancel := chromedp.NewExecAllocator(ctx, chromedp.NoSandbox, chromedp.NoFirstRun, chromedp.NoDefaultBrowserCheck, chromedp.Headless, chromedp.DisableGPU, ) defer cancel() ctx, cancel = chromedp.NewContext(allocCtx, chromedp.WithLogf(func(format string, v ...interface{}) { if opts.Logger != nil { opts.Logger.Debug(fmt.Sprintf(format, v...)) } })) defer cancel() // 设置超时 ctx, cancel = context.WithTimeout(ctx, opts.Timeout) defer cancel() var buf []byte // 执行截图任务 err := chromedp.Run(ctx, chromedp.Navigate(url), chromedp.WaitVisible(selector, chromedp.ByQuery), chromedp.Sleep(opts.WaitTime), chromedp.Screenshot(selector, &buf, chromedp.NodeVisible), ) if err != nil { return "", fmt.Errorf("failed to capture element screenshot: %w", err) } // 转换为base64 base64Str := base64.StdEncoding.EncodeToString(buf) return base64Str, nil } // ScreenshotHTMLToMessageChain 对HTML内容进行截图并返回包含图片的消息链 func ScreenshotHTMLToMessageChain(ctx context.Context, htmlContent string, opts *ScreenshotOptions) (protocol.MessageChain, error) { base64Data, err := ScreenshotHTML(ctx, htmlContent, opts) if err != nil { return nil, err } chain := protocol.NewMessageChain() chain = chain.AppendImageFromBase64(base64Data) return chain, nil } // ScreenshotURLToMessageChain 对URL进行截图并返回包含图片的消息链 func ScreenshotURLToMessageChain(ctx context.Context, url string, opts *ScreenshotOptions) (protocol.MessageChain, error) { base64Data, err := ScreenshotURL(ctx, url, opts) if err != nil { return nil, err } chain := protocol.NewMessageChain() chain = chain.AppendImageFromBase64(base64Data) return chain, nil } // ============================================================================ // HTML 模板功能 // ============================================================================ // RenderHTMLTemplate 渲染HTML模板并返回HTML字符串 func RenderHTMLTemplate(tmplContent string, data interface{}) (string, error) { tmpl, err := template.New("html").Parse(tmplContent) if err != nil { return "", fmt.Errorf("failed to parse template: %w", err) } var buf bytes.Buffer if err := tmpl.Execute(&buf, data); err != nil { return "", fmt.Errorf("failed to execute template: %w", err) } return buf.String(), nil } // RenderHTMLTemplateFromFile 从文件加载并渲染HTML模板 func RenderHTMLTemplateFromFile(tmplPath string, data interface{}) (string, error) { content, err := os.ReadFile(tmplPath) if err != nil { return "", fmt.Errorf("failed to read template file: %w", err) } tmpl, err := template.New(filepath.Base(tmplPath)).Parse(string(content)) if err != nil { return "", fmt.Errorf("failed to parse template: %w", err) } var buf bytes.Buffer if err := tmpl.Execute(&buf, data); err != nil { return "", fmt.Errorf("failed to execute template: %w", err) } return buf.String(), nil } // ScreenshotTemplate 渲染HTML模板并截图,返回base64编码 func ScreenshotTemplate(ctx context.Context, tmplContent string, data interface{}, opts *ScreenshotOptions) (string, error) { html, err := RenderHTMLTemplate(tmplContent, data) if err != nil { return "", err } return ScreenshotHTML(ctx, html, opts) } // ScreenshotTemplateFromFile 从文件加载模板,渲染并截图,返回base64编码 func ScreenshotTemplateFromFile(ctx context.Context, tmplPath string, data interface{}, opts *ScreenshotOptions) (string, error) { html, err := RenderHTMLTemplateFromFile(tmplPath, data) if err != nil { return "", err } return ScreenshotHTML(ctx, html, opts) } // ScreenshotTemplateToMessageChain 渲染HTML模板并截图,返回包含图片的消息链 func ScreenshotTemplateToMessageChain(ctx context.Context, tmplContent string, data interface{}, opts *ScreenshotOptions) (protocol.MessageChain, error) { base64Data, err := ScreenshotTemplate(ctx, tmplContent, data, opts) if err != nil { return nil, err } chain := protocol.NewMessageChain() chain = chain.AppendImageFromBase64(base64Data) return chain, nil } // ScreenshotTemplateFromFileToMessageChain 从文件加载模板,渲染并截图,返回包含图片的消息链 func ScreenshotTemplateFromFileToMessageChain(ctx context.Context, tmplPath string, data interface{}, opts *ScreenshotOptions) (protocol.MessageChain, error) { base64Data, err := ScreenshotTemplateFromFile(ctx, tmplPath, data, opts) if err != nil { return nil, err } chain := protocol.NewMessageChain() chain = chain.AppendImageFromBase64(base64Data) return chain, nil }