cellbot/pkg/utils/screenshot.go

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
}