FranklinWH-Toolbox/service/services/dataDownloadRenderService.go

package services

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/csv"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html"
	"io"
	"strings"
	"time"

	"fwh-toolbox-server/cache"
	"fwh-toolbox-server/domain"
	minioRepo "fwh-toolbox-server/repositories/minio"
	mongoRepo "fwh-toolbox-server/repositories/mongo"
	pgsqlRepo "fwh-toolbox-server/repositories/pgsql"

	"github.com/xuri/excelize/v2"
	"go.mongodb.org/mongo-driver/mongo"
	"go.uber.org/zap"
)

const renderCacheTTL = 72 * time.Hour

func ReadRenderedDataDownloadDocument(ctx context.Context, documentID, format string) (*domain.DataDownloadArtifact, []byte, error) {
	document, format, lock, err := lockRenderedDataDownloadDocument(ctx, documentID, format)
	if err != nil {
		return nil, nil, err
	}
	defer lock.SafeRelease(context.Background())
	artifact, data, err := resolveRenderedDataDownloadDocument(ctx, document, format)
	if err != nil || data != nil {
		return artifact, data, err
	}
	data, err = readDataDownloadArtifact(ctx, artifact)
	return artifact, data, err
}

func OpenRenderedDataDownloadDocument(ctx context.Context, documentID, format string) (*domain.DataDownloadArtifact, io.ReadCloser, error) {
	document, format, lock, err := lockRenderedDataDownloadDocument(ctx, documentID, format)
	if err != nil {
		return nil, nil, err
	}
	artifact, data, err := resolveRenderedDataDownloadDocument(ctx, document, format)
	if err != nil {
		lock.SafeRelease(context.Background())
		return nil, nil, err
	}
	var reader io.ReadCloser
	if data != nil {
		reader = io.NopCloser(bytes.NewReader(data))
	} else {
		reader, err = minioRepo.Open(ctx, artifact.ObjectKey)
		if err != nil {
			lock.SafeRelease(context.Background())
			return nil, nil, err
		}
	}
	return artifact, newDataDownloadContentReader(reader, lock), nil
}

func lockRenderedDataDownloadDocument(ctx context.Context, documentID, format string) (*domain.DataDownloadDocument, string, *cache.RedisLock, error) {
	document, err := mongoRepo.GetDataDownloadDocument(ctx, documentID)
	if err != nil {
		return nil, "", nil, err
	}
	format = strings.ToLower(strings.TrimPrefix(format, "."))
	if !supportedRenderFormat(document.Kind, format) {
		return nil, "", nil, fmt.Errorf("format %s is not supported for %s", format, document.Kind)
	}

	// This lock is shared by writers. It is deliberately exclusive at Redis level;
	// PostgreSQL's advisory shared lock still permits safe concurrent database readers.
	lock, err := cache.Lock(ctx, "data-download:content:"+documentID, 10*time.Minute, 200*time.Millisecond, 30*time.Second)
	if err != nil {
		return nil, "", nil, err
	}
	document, err = mongoRepo.GetDataDownloadDocument(ctx, documentID)
	if err != nil {
		lock.SafeRelease(context.Background())
		return nil, "", nil, err
	}
	return document, format, lock, nil
}

func resolveRenderedDataDownloadDocument(ctx context.Context, document *domain.DataDownloadDocument, format string) (*domain.DataDownloadArtifact, []byte, error) {
	if cached, err := currentArtifact(ctx, document, format); err == nil {
		zap.L().Debug("[DataDownload] 渲染缓存命中",
			zap.String("documentId", document.ID), zap.String("format", format),
			zap.Int64("generation", cached.ContentGeneration))
		return cached, nil, nil
	}
	state, rows, err := pgsqlRepo.ReadDataDownloadRows(document.ID, document.Kind)
	if err != nil {
		return nil, nil, err
	}
	data, contentType, err := encodeRenderedRows(document, rows, format)
	if err != nil {
		return nil, nil, err
	}
	objectKey, err := minioRepo.PutRenderArtifact(ctx, document.ID, state.ContentGeneration, format, data, contentType)
	if err != nil {
		return nil, nil, err
	}
	sum := sha256.Sum256(data)
	now := time.Now().UTC()
	artifact := &domain.DataDownloadArtifact{DocumentID: document.ID, Format: format, ContentGeneration: state.ContentGeneration, ObjectKey: objectKey, Filename: renderedFilename(document, format), Size: int64(len(data)), Checksum: hex.EncodeToString(sum[:]), ExpiresAt: now.Add(renderCacheTTL), CreatedAt: now}
	if err := mongoRepo.PublishDataDownloadArtifact(ctx, artifact); err != nil {
		_ = minioRepo.Delete(ctx, objectKey)
		return nil, nil, err
	}
	zap.L().Debug("[DataDownload] 渲染完成并发布 artifact",
		zap.String("documentId", document.ID), zap.String("format", format),
		zap.Int64("generation", state.ContentGeneration), zap.Int("bytes", len(data)))
	return artifact, data, nil
}

type dataDownloadContentReader struct {
	io.ReadCloser
	lock   *cache.RedisLock
	cancel context.CancelFunc
	done   chan struct{}
}

func newDataDownloadContentReader(reader io.ReadCloser, lock *cache.RedisLock) *dataDownloadContentReader {
	ctx, cancel := context.WithCancel(context.Background())
	result := &dataDownloadContentReader{ReadCloser: reader, lock: lock, cancel: cancel, done: make(chan struct{})}
	go func() {
		defer close(result.done)
		ticker := time.NewTicker(3 * time.Minute)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				renewCtx, renewCancel := context.WithTimeout(context.Background(), 10*time.Second)
				_ = lock.Extend(renewCtx, 10*time.Minute)
				renewCancel()
			}
		}
	}()
	return result
}

func (r *dataDownloadContentReader) Close() error {
	r.cancel()
	<-r.done
	err := r.ReadCloser.Close()
	r.lock.SafeRelease(context.Background())
	return err
}

func readDataDownloadArtifact(ctx context.Context, artifact *domain.DataDownloadArtifact) ([]byte, error) {
	reader, err := minioRepo.Open(ctx, artifact.ObjectKey)
	if err != nil {
		return nil, err
	}
	defer reader.Close()
	return io.ReadAll(reader)
}

func currentArtifact(ctx context.Context, document *domain.DataDownloadDocument, format string) (*domain.DataDownloadArtifact, error) {
	slot, err := mongoRepo.GetDataDownloadRenderSlot(ctx, document.ID, format)
	if err != nil {
		return nil, err
	}
	if slot.ArtifactID == nil || slot.ContentGeneration != document.ContentGeneration {
		return nil, mongo.ErrNoDocuments
	}
	artifact, err := mongoRepo.GetDataDownloadArtifact(ctx, *slot.ArtifactID)
	if err != nil {
		return nil, err
	}
	if !artifact.ExpiresAt.After(time.Now()) {
		return nil, mongo.ErrNoDocuments
	}
	return artifact, nil
}

func supportedRenderFormat(kind domain.DataDownloadKind, format string) bool {
	switch kind {
	case domain.DataDownloadPowerFlow:
		return format == "html" || format == "json"
	case domain.DataDownloadMqttLog, domain.DataDownloadEventAlarm:
		return format == "xlsx" || format == "csv" || format == "json"
	case domain.DataDownloadUpdateLog:
		return format == "xlsx" || format == "json"
	default:
		return false
	}
}

func encodeRenderedRows(document *domain.DataDownloadDocument, rows []json.RawMessage, format string) ([]byte, string, error) {
	switch format {
	case "json":
		data, err := json.MarshalIndent(rows, "", "  ")
		return data, "application/json", err
	case "csv":
		var out bytes.Buffer
		writer := csv.NewWriter(&out)
		headers, table := dataDownloadTable(document.Kind, rows)
		_ = writer.Write(headers)
		for _, row := range table {
			_ = writer.Write(row)
		}
		writer.Flush()
		return out.Bytes(), "text/csv; charset=utf-8", writer.Error()
	case "html":
		var out strings.Builder
		out.WriteString("<!doctype html><html><head><meta charset=\"utf-8\"><title>PowerFlow</title></head><body><h1>PowerFlow ")
		out.WriteString(html.EscapeString(document.SiteSN + " " + document.Date))
		out.WriteString("</h1><pre>")
		pretty, _ := json.MarshalIndent(rows, "", "  ")
		out.WriteString(html.EscapeString(string(pretty)))
		out.WriteString("</pre></body></html>")
		return []byte(out.String()), "text/html; charset=utf-8", nil
	case "xlsx":
		file := excelize.NewFile()
		sheet := file.GetSheetName(0)
		headers, table := dataDownloadTable(document.Kind, rows)
		for column, header := range headers {
			cell, _ := excelize.CoordinatesToCellName(column+1, 1)
			_ = file.SetCellValue(sheet, cell, header)
		}
		for rowIndex, row := range table {
			for column, value := range row {
				cell, _ := excelize.CoordinatesToCellName(column+1, rowIndex+2)
				_ = file.SetCellValue(sheet, cell, value)
			}
		}
		buffer, err := file.WriteToBuffer()
		_ = file.Close()
		if err != nil {
			return nil, "", err
		}
		return buffer.Bytes(), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", nil
	default:
		return nil, "", fmt.Errorf("unsupported format")
	}
}

func dataDownloadTable(kind domain.DataDownloadKind, rows []json.RawMessage) ([]string, [][]string) {
	fields := []string{"payload"}
	headers := []string{"Payload"}
	switch kind {
	case domain.DataDownloadMqttLog:
		fields, headers = []string{"cmdType", "logTime", "logData"}, []string{"Type", "Time", "Data"}
	case domain.DataDownloadEventAlarm:
		fields = []string{"id", "logTime", "logType", "alarmCode", "logName", "startTime", "endTime", "logData"}
		headers = []string{"ID", "Time", "Type", "AlarmCode", "Name", "StartTime", "EndTime", "Data"}
	}
	table := make([][]string, 0, len(rows))
	for index, raw := range rows {
		if len(fields) == 1 {
			table = append(table, []string{string(raw)})
			continue
		}
		var object map[string]any
		if err := json.Unmarshal(raw, &object); err != nil {
			table = append(table, []string{string(raw)})
			continue
		}
		row := make([]string, len(fields))
		for column, field := range fields {
			value := object[field]
			if value == nil && kind == domain.DataDownloadEventAlarm && field == "id" {
				value = index + 1
			}
			if value == nil && kind == domain.DataDownloadEventAlarm && field == "endTime" {
				row[column] = "null"
				continue
			}
			switch typed := value.(type) {
			case string:
				row[column] = typed
			case nil:
				row[column] = ""
			default:
				encoded, _ := json.Marshal(typed)
				row[column] = string(encoded)
			}
		}
		table = append(table, row)
	}
	return headers, table
}

func renderedFilename(document *domain.DataDownloadDocument, format string) string {
	parts := []string{document.SiteSN, string(document.Kind)}
	if document.Date != "" {
		parts = append(parts, document.Date)
	}
	if document.Subtype != "" {
		parts = append(parts, document.Subtype)
	}
	return strings.Join(parts, "_") + "." + format
}