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

package pdf

import (
	"context"
	"net/url"
	"strings"
	"sync"
	"time"

	"go.uber.org/zap"
)

// QYGLReportPDFStatus 企业报告 PDF 预生成状态（供前端轮询）
type QYGLReportPDFStatus string

const (
	QYGLReportPDFStatusNone       QYGLReportPDFStatus = "none"
	QYGLReportPDFStatusPending    QYGLReportPDFStatus = "pending"
	QYGLReportPDFStatusGenerating QYGLReportPDFStatus = "generating"
	QYGLReportPDFStatusReady      QYGLReportPDFStatus = "ready"
	QYGLReportPDFStatusFailed     QYGLReportPDFStatus = "failed"
)

// QYGLReportPDFPregen 异步预渲染企业报告 PDF（headless Chrome 访问公网可访问的报告 HTML）
type QYGLReportPDFPregen struct {
	logger  *zap.Logger
	cache   *PDFCacheManager
	baseURL string // 已 trim，无尾斜杠；空则禁用

	mu     sync.RWMutex
	states map[string]*qyglPDFState
}

type qyglPDFState struct {
	Status    QYGLReportPDFStatus
	Message   string
	UpdatedAt time.Time
}

// NewQYGLReportPDFPregen baseURL 须为外部可访问基址（如 https://api.example.com），为空时不预生成
func NewQYGLReportPDFPregen(logger *zap.Logger, cache *PDFCacheManager, baseURL string) *QYGLReportPDFPregen {
	if logger == nil {
		logger = zap.NewNop()
	}
	u := strings.TrimSpace(baseURL)
	u = strings.TrimRight(u, "/")
	return &QYGLReportPDFPregen{
		logger:  logger,
		cache:   cache,
		baseURL: u,
		states:  make(map[string]*qyglPDFState),
	}
}

// ScheduleQYGLReportPDF 实现 processors.QYGLReportPDFScheduler
func (p *QYGLReportPDFPregen) ScheduleQYGLReportPDF(ctx context.Context, reportID string) {
	p.schedule(ctx, reportID)
}

func (p *QYGLReportPDFPregen) schedule(_ context.Context, reportID string) {
	if p == nil || p.baseURL == "" || reportID == "" || p.cache == nil {
		return
	}
	if b, hit, _, err := p.cache.GetByReportID(reportID); hit && err == nil && len(b) > 0 {
		p.setState(reportID, QYGLReportPDFStatusReady, "")
		return
	}

	p.mu.Lock()
	if st, ok := p.states[reportID]; ok {
		if st.Status == QYGLReportPDFStatusGenerating || st.Status == QYGLReportPDFStatusPending {
			p.mu.Unlock()
			return
		}
		if st.Status == QYGLReportPDFStatusReady {
			p.mu.Unlock()
			return
		}
	}
	p.states[reportID] = &qyglPDFState{
		Status:    QYGLReportPDFStatusPending,
		Message:   "",
		UpdatedAt: time.Now(),
	}
	p.mu.Unlock()

	go p.runGeneration(reportID)
}

func (p *QYGLReportPDFPregen) runGeneration(reportID string) {
	p.setState(reportID, QYGLReportPDFStatusGenerating, "")

	fullURL := p.baseURL + "/reports/qygl/" + url.PathEscape(reportID)
	gen := NewHTMLPDFGenerator(p.logger)

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

	pdfBytes, err := gen.GenerateFromURL(ctx, fullURL)
	if err != nil {
		p.logger.Error("企业报告 PDF 预生成失败",
			zap.String("report_id", reportID),
			zap.String("url", fullURL),
			zap.Error(err),
		)
		p.setState(reportID, QYGLReportPDFStatusFailed, "生成失败："+err.Error())
		return
	}
	if len(pdfBytes) == 0 {
		p.setState(reportID, QYGLReportPDFStatusFailed, "生成的 PDF 为空")
		return
	}
	if err := p.cache.SetByReportID(reportID, pdfBytes); err != nil {
		p.logger.Error("企业报告 PDF 写入缓存失败",
			zap.String("report_id", reportID),
			zap.Error(err),
		)
		p.setState(reportID, QYGLReportPDFStatusFailed, "写入缓存失败")
		return
	}
	p.setState(reportID, QYGLReportPDFStatusReady, "")
	p.logger.Info("企业报告 PDF 预生成完成",
		zap.String("report_id", reportID),
		zap.Int("size", len(pdfBytes)),
	)
}

func (p *QYGLReportPDFPregen) setState(reportID string, st QYGLReportPDFStatus, msg string) {
	if p == nil {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.states == nil {
		p.states = make(map[string]*qyglPDFState)
	}
	p.states[reportID] = &qyglPDFState{
		Status:    st,
		Message:   msg,
		UpdatedAt: time.Now(),
	}
}

// Status 返回当前状态；若磁盘缓存已命中则始终为 ready
func (p *QYGLReportPDFPregen) Status(reportID string) (QYGLReportPDFStatus, string) {
	if p == nil || reportID == "" {
		return QYGLReportPDFStatusNone, ""
	}
	if p.cache != nil {
		if b, hit, _, err := p.cache.GetByReportID(reportID); hit && err == nil && len(b) > 0 {
			return QYGLReportPDFStatusReady, ""
		}
	}
	p.mu.RLock()
	defer p.mu.RUnlock()
	st := p.states[reportID]
	if st == nil {
		if p.baseURL == "" {
			return QYGLReportPDFStatusNone, "未启用预生成"
		}
		return QYGLReportPDFStatusNone, ""
	}
	return st.Status, st.Message
}

// MarkReadyAfterUpload 同步生成成功后与预生成状态对齐（可选）
func (p *QYGLReportPDFPregen) MarkReadyAfterUpload(reportID string) {
	if p == nil || reportID == "" {
		return
	}
	p.setState(reportID, QYGLReportPDFStatusReady, "")
}

// Enabled 是否配置了预生成基址
func (p *QYGLReportPDFPregen) Enabled() bool {
	return p != nil && p.baseURL != ""
}