hyapi-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"
)

// A4 纵向 @ 96dpi，与 CSS @page / 报告页 max-width 对齐
const (
	htmlPDFViewportWidth  = 794
	htmlPDFViewportHeight = 1123
)

// 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 视口渲染 → 等待报告 DOM → 滚动触发展开 → 切换 print 媒体 → PrintToPDF（边距由 CSS @page 控制）
func (g *HTMLPDFGenerator) GenerateFromURL(ctx context.Context, url string) ([]byte, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	timeoutCtx, cancel := context.WithTimeout(ctx, 90*time.Second)
	defer cancel()

	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.Flag("headless", true),
		chromedp.Flag("disable-gpu", true),
		chromedp.Flag("no-sandbox", true),
		chromedp.Flag("disable-dev-shm-usage", true),
	)
	allocCtx, allocCancel := chromedp.NewExecAllocator(timeoutCtx, opts...)
	defer allocCancel()

	chromeCtx, cancelChrome := chromedp.NewContext(allocCtx)
	defer cancelChrome()

	var pdfBuf []byte

	tasks := chromedp.Tasks{
		chromedp.EmulateViewport(htmlPDFViewportWidth, htmlPDFViewportHeight),
		chromedp.Navigate(url),
		chromedp.WaitReady("body", chromedp.ByQuery),
		chromedp.WaitVisible(".page", chromedp.ByQuery),
		waitReportReadyForPDF(),
		waitReportPosterImages(),
		scrollPageForPrintLayout(),
		chromedp.ActionFunc(func(ctx context.Context) error {
			return chromedp.Evaluate(`document.body.classList.add('pdf-export')`, nil).Do(ctx)
		}),
		chromedp.ActionFunc(func(ctx context.Context) error {
			time.Sleep(150 * time.Millisecond)
			return nil
		}),
		chromedp.ActionFunc(func(ctx context.Context) error {
			g.logger.Info("开始通过 headless Chrome 生成企业报告 PDF", zap.String("url", url))
			buf, _, err := page.PrintToPDF().
				WithPrintBackground(true).
				WithPaperWidth(8.27).
				WithPaperHeight(11.69).
				WithMarginTop(0.39).
				WithMarginBottom(0.39).
				WithMarginLeft(0.47).
				WithMarginRight(0.47).
				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
}

// waitReportReadyForPDF 等待企业报告脚本渲染完成（qiye.html 设置 data-qygl-report-ready）
func waitReportReadyForPDF() chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		deadline := time.Now().Add(45 * time.Second)
		for {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			var ready string
			_ = chromedp.Evaluate(`document.body && document.body.dataset.qyglReportReady || ''`, &ready).Do(ctx)
			if ready == "1" {
				return nil
			}
			var sectionCount float64
			_ = chromedp.Evaluate(`document.querySelectorAll('#reportSections .section-card, .report-sections .section-block').length`, &sectionCount).Do(ctx)
			if sectionCount > 0 {
				return nil
			}
			if time.Now().After(deadline) {
				return fmt.Errorf("等待报告内容渲染超时")
			}
			time.Sleep(200 * time.Millisecond)
		}
	})
}

func waitReportPosterImages() chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		deadline := time.Now().Add(20 * time.Second)
		for {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			var pending float64
			js := `(function(){
				var imgs = document.querySelectorAll('img.header-poster, img.footer-poster');
				if (!imgs.length) return 0;
				for (var i = 0; i < imgs.length; i++) {
					var img = imgs[i];
					if (!img.complete || img.naturalWidth === 0) return 1;
				}
				return 0;
			})()`
			if err := chromedp.Evaluate(js, &pending).Do(ctx); err != nil {
				return err
			}
			if pending == 0 {
				return nil
			}
			if time.Now().After(deadline) {
				return fmt.Errorf("等待报告海报图片加载超时")
			}
			time.Sleep(150 * time.Millisecond)
		}
	})
}

func scrollPageForPrintLayout() chromedp.Action {
	return chromedp.ActionFunc(func(ctx context.Context) error {
		var height float64
		if err := chromedp.Evaluate(`Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)`, &height).Do(ctx); err != nil {
			return err
		}
		step := float64(900)
		if height < step {
			step = height
		}
		if step <= 0 {
			return nil
		}
		for y := float64(0); y <= height; y += step {
			js := fmt.Sprintf(`window.scrollTo(0, %f);`, y)
			if err := chromedp.Evaluate(js, nil).Do(ctx); err != nil {
				return err
			}
			time.Sleep(120 * time.Millisecond)
		}
		if err := chromedp.Evaluate(`window.scrollTo(0, 0);`, nil).Do(ctx); err != nil {
			return err
		}
		time.Sleep(300 * time.Millisecond)
		return nil
	})
}