FranklinWH-Toolbox/service/services/dataDownloadPackageService.go

package services

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"path"
	"strings"
	"time"

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

	"go.uber.org/zap"
)

func BuildDataDownloadPackage(ctx context.Context, job *domain.DataDownloadPackageJob, progress func(done, total int64) error) ([]byte, int64, error) {
	type input struct {
		taskID, logicalPath, objectKey, filename string
		data                                     []byte
		document                                 *domain.DataDownloadDocument
		artifact                                 *domain.DataDownloadArtifact
	}
	var inputs []input
	type reportFailure struct {
		TaskID        string `json:"taskId"`
		ItemID        string `json:"itemId"`
		DocumentID    string `json:"documentId"`
		SourceFileID  string `json:"sourceFileId"`
		FileName      string `json:"fileName,omitempty"`
		FailureStage  string `json:"failureStage"`
		FailureReason string `json:"failureReason"`
		AttemptCount  int    `json:"attemptCount"`
		DeviceSN      string `json:"deviceSn,omitempty"`
		DeviceType    string `json:"deviceType,omitempty"`
		Date          string `json:"date,omitempty"`
	}
	type reportTask struct {
		TaskID    string                    `json:"taskId"`
		Status    domain.DataDownloadStatus `json:"status"`
		Succeeded int64                     `json:"succeeded"`
		Failed    int64                     `json:"failed"`
	}
	var failures []reportFailure
	var reportTasks []reportTask
	for _, taskID := range job.TaskIDs {
		task, err := mongoRepo.GetDataDownloadTask(ctx, taskID)
		if err != nil {
			return nil, 0, err
		}
		reportTasks = append(reportTasks, reportTask{TaskID: taskID.Hex(), Status: task.Status, Succeeded: task.Succeeded, Failed: task.Failed})
		items, err := mongoRepo.ListDataDownloadItems(ctx, taskID)
		if err != nil {
			return nil, 0, err
		}
		itemByID := make(map[string]domain.DataDownloadItem, len(items))
		for i := range items {
			itemByID[items[i].ID.Hex()] = items[i]
		}
		attempts, err := mongoRepo.ListDataDownloadFileAttempts(ctx, taskID)
		if err != nil {
			return nil, 0, err
		}
		for i := range attempts {
			if attempts[i].Status != domain.DataDownloadFailed {
				continue
			}
			item := itemByID[attempts[i].ItemID.Hex()]
			failures = append(failures, reportFailure{
				TaskID: taskID.Hex(), ItemID: attempts[i].ItemID.Hex(), DocumentID: attempts[i].DocumentID,
				SourceFileID: attempts[i].SourceFileID, FileName: attempts[i].FileName,
				FailureStage: attempts[i].FailureStage, FailureReason: attempts[i].FailureReason,
				AttemptCount: attempts[i].AttemptCount, DeviceSN: item.DeviceSN, DeviceType: item.DeviceType, Date: item.Date,
			})
		}
		links, err := mongoRepo.ListDataDownloadTaskDocuments(ctx, taskID)
		if err != nil {
			return nil, 0, err
		}
		for _, link := range links {
			if link.Empty {
				continue
			}
			document, err := mongoRepo.GetDataDownloadDocument(ctx, link.DocumentID)
			if err != nil {
				return nil, 0, err
			}
			if document.Kind == domain.DataDownloadFaultRecord {
				lock, err := cache.Lock(ctx, "data-download:content:"+document.ID, 10*time.Minute, 100*time.Millisecond, 30*time.Second)
				if err != nil {
					return nil, 0, err
				}
				entries, err := mongoRepo.ListDataDownloadFaultRecordEntries(ctx, document.ID)
				if err != nil {
					lock.SafeRelease(context.Background())
					return nil, 0, err
				}
				for i := range entries {
					reader, openErr := minioRepo.Open(ctx, entries[i].ObjectKey)
					if openErr != nil {
						lock.SafeRelease(context.Background())
						return nil, 0, openErr
					}
					data, readErr := io.ReadAll(reader)
					_ = reader.Close()
					if readErr != nil {
						lock.SafeRelease(context.Background())
						return nil, 0, readErr
					}
					filename := fmt.Sprintf("%d-%d-%s", entries[i].FaultID, entries[i].SubRecordID, entries[i].FileName)
					inputs = append(inputs, input{taskID: taskID.Hex(), logicalPath: link.LogicalPath, objectKey: entries[i].ObjectKey, filename: filename, data: data, document: document})
				}
				lock.SafeRelease(context.Background())
				continue
			}
			entry := input{taskID: taskID.Hex(), logicalPath: link.LogicalPath, document: document}
			if document.BinaryObjectKey != "" {
				lock, err := cache.Lock(ctx, "data-download:content:"+document.ID, 10*time.Minute, 100*time.Millisecond, 30*time.Second)
				if err != nil {
					return nil, 0, err
				}
				current, err := mongoRepo.GetDataDownloadDocument(ctx, document.ID)
				if err != nil {
					lock.SafeRelease(context.Background())
					return nil, 0, err
				}
				reader, err := minioRepo.Open(ctx, current.BinaryObjectKey)
				if err != nil {
					lock.SafeRelease(context.Background())
					return nil, 0, err
				}
				entry.data, err = io.ReadAll(reader)
				_ = reader.Close()
				lock.SafeRelease(context.Background())
				if err != nil {
					return nil, 0, err
				}
				entry.objectKey = current.BinaryObjectKey
			} else {
				artifact, data, err := ReadRenderedDataDownloadDocument(ctx, document.ID, defaultDataDownloadFormat(document.Kind))
				if err != nil {
					return nil, 0, err
				}
				entry.data = data
				entry.artifact, entry.objectKey = artifact, artifact.ObjectKey
			}
			inputs = append(inputs, entry)
		}
	}
	if len(inputs) == 0 && len(failures) == 0 {
		return nil, 0, fmt.Errorf("no readable non-empty documents")
	}
	if len(failures) > 0 {
		report, err := json.MarshalIndent(struct {
			GeneratedAt time.Time       `json:"generatedAt"`
			Tasks       []reportTask    `json:"tasks"`
			Failures    []reportFailure `json:"failures"`
		}{GeneratedAt: time.Now().UTC(), Tasks: reportTasks, Failures: failures}, "", "  ")
		if err != nil {
			return nil, 0, err
		}
		inputs = append(inputs, input{logicalPath: "", filename: "DataDownloadReport.json", data: report})
	}
	if progress != nil {
		if err := progress(0, int64(len(inputs))); err != nil {
			return nil, 0, err
		}
	}
	zap.L().Debug("[DataDownload] 打包收集完成，开始写 ZIP",
		zap.Int("fileCount", len(inputs)),
		zap.Int("failureCount", len(failures)),
		zap.Int("taskCount", len(job.TaskIDs)),
	)
	var buffer bytes.Buffer
	archive := zip.NewWriter(&buffer)
	for i, entry := range inputs {
		var reader io.ReadCloser
		if entry.data != nil {
			reader = io.NopCloser(bytes.NewReader(entry.data))
		} else {
			opened, err := minioRepo.Open(ctx, entry.objectKey)
			if err != nil {
				zap.L().Warn("[DataDownload] 打包读取对象失败", zap.String("objectKey", entry.objectKey), zap.Int("index", i), zap.String("err", err.Error()))
				_ = archive.Close()
				return nil, int64(i), err
			}
			reader = opened
		}
		filename := entry.filename
		if filename == "" {
			filename = path.Base(entry.objectKey)
		}
		if entry.artifact != nil {
			filename = entry.artifact.Filename
		}
		logical := cleanZipPath(entry.logicalPath)
		if len(job.TaskIDs) > 1 {
			logical = path.Join(entry.taskID, logical)
		}
		writer, err := archive.Create(path.Join(logical, filename))
		if err == nil {
			_, err = io.Copy(writer, reader)
		}
		_ = reader.Close()
		if err != nil {
			_ = archive.Close()
			return nil, int64(i), err
		}
		if progress != nil {
			if err := progress(int64(i+1), int64(len(inputs))); err != nil {
				_ = archive.Close()
				return nil, int64(i + 1), err
			}
		}
	}
	if err := archive.Close(); err != nil {
		return nil, 0, err
	}
	return buffer.Bytes(), int64(len(inputs)), nil
}

func defaultDataDownloadFormat(kind domain.DataDownloadKind) string {
	switch kind {
	case domain.DataDownloadPowerFlow:
		return "html"
	case domain.DataDownloadMqttLog, domain.DataDownloadEventAlarm, domain.DataDownloadUpdateLog:
		return "xlsx"
	default:
		return ""
	}
}

func cleanZipPath(value string) string {
	value = strings.ReplaceAll(value, "\\", "/")
	value = strings.TrimPrefix(path.Clean("/"+value), "/")
	if value == "." {
		return "documents"
	}
	return value
}