tyapi-server/internal/shared/pdf/html_pdf_generator.go

package pdf

import (
	"context"
	"fmt"
	"time"

	"github.com/chromedp/cdproto/page"
	"github.com/chromedp/chromedp"
	"go.uber.org/zap"
)

// HTMLPDFGenerator 使用 headless Chrome 将 HTML 页面渲染为 PDF
type HTMLPDFGenerator struct {
	logger *zap.Logger
}

// NewHTMLPDFGenerator 创建 HTMLPDFGenerator
func NewHTMLPDFGenerator(logger *zap.Logger) *HTMLPDFGenerator {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &HTMLPDFGenerator{
		logger: logger,
	}
}

// GenerateFromURL 使用 headless Chrome 打开指定 URL，并导出为 PDF 字节流
// 这里固定使用 A4 纵向纸张，开启背景打印
func (g *HTMLPDFGenerator) GenerateFromURL(ctx context.Context, url string) ([]byte, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	// 整个生成过程增加超时时间，避免长时间卡死
	timeoutCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	// 创建 Chrome 上下文（使用系统默认的 headless Chrome/Chromium）
	chromeCtx, cancelChrome := chromedp.NewContext(timeoutCtx)
	defer cancelChrome()

	var pdfBuf []byte

	tasks := chromedp.Tasks{
		chromedp.Navigate(url),
		// 等待页面主体和报告容器就绪，确保数据渲染完成
		chromedp.WaitReady("body", chromedp.ByQuery),
		chromedp.WaitVisible(".page", chromedp.ByQuery),
		chromedp.ActionFunc(func(ctx context.Context) error {
			g.logger.Info("开始通过 headless Chrome 生成企业报告 PDF", zap.String("url", url))
			var (
				buf []byte
				err error
			)
			buf, _, err = page.PrintToPDF().
				WithPrintBackground(true).
				WithPaperWidth(8.27).  // A4 宽度（英寸 -> 约 210mm）
				WithPaperHeight(11.69). // A4 高度（英寸 -> 约 297mm）
				WithMarginTop(0.4).
				WithMarginBottom(0.4).
				WithMarginLeft(0.4).
				WithMarginRight(0.4).
				Do(ctx)
			if err == nil {
				pdfBuf = buf
			}
			return err
		}),
	}

	if err := chromedp.Run(chromeCtx, tasks); err != nil {
		g.logger.Error("使用 headless Chrome 生成 HTML 报告 PDF 失败", zap.String("url", url), zap.Error(err))
		return nil, err
	}

	if len(pdfBuf) == 0 {
		return nil, fmt.Errorf("生成的 PDF 内容为空")
	}

	g.logger.Info("通过 headless Chrome 生成企业报告 PDF 成功",
		zap.String("url", url),
		zap.Int("pdf_size", len(pdfBuf)),
	)

	return pdfBuf, nil
}