FranklinWH-Toolbox/service/services/dataDownloadBinaryService.go

package services

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"path"
	"strconv"
	"time"

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

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

func ExecuteBinaryDataDownloadItem(ctx context.Context, item *domain.DataDownloadItem, awaitRequestIDs ...string) error {
	awaitRequestID := ""
	if len(awaitRequestIDs) > 0 {
		awaitRequestID = awaitRequestIDs[0]
	}
	if item.Kind == domain.DataDownloadFaultRecord {
		return executeFaultRecordDataDownloadItem(ctx, item, awaitRequestID)
	}
	ready, lease, err := prepareDataDownloadRequest(ctx, item, awaitRequestID)
	if err != nil {
		return err
	}
	if ready != nil {
		return linkDataDownloadItem(ctx, item, ready, ready.Empty)
	}
	defer lease.abort()
	documentID := lease.documentID
	fetchedAt := time.Now().UTC()
	var data []byte
	var filename, contentType string
	switch item.Kind {
	case domain.DataDownloadAlarmSnapshot:
		values, err := jobs.GetAlarmSnapshotData(ctx, item.SiteSN, item.Date)
		if err != nil {
			zap.L().Warn("[DataDownload] AlarmSnapshot 上游查询失败", zap.String("sn", item.SiteSN), zap.String("date", item.Date), zap.String("err", err.Error()))
			return err
		}
		zap.L().Debug("[DataDownload] AlarmSnapshot 上游返回", zap.String("sn", item.SiteSN), zap.String("date", item.Date), zap.Int("rows", len(values)))
		if len(values) > 0 {
			file := excelize.NewFile()
			sheet := file.GetSheetName(0)
			_ = file.SetSheetRow(sheet, "A1", &[]any{"ID", "Time", "Data"})
			for i, value := range values {
				_ = file.SetSheetRow(sheet, fmt.Sprintf("A%d", i+2), &[]any{value.ID, value.LogTime.Format(time.RFC3339), value.JSONData})
			}
			buffer, err := file.WriteToBuffer()
			_ = file.Close()
			if err != nil {
				return err
			}
			data = buffer.Bytes()
		}
		filename, contentType = fmt.Sprintf("AlarmSnapshot_%s.xlsx", item.Date), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
	case domain.DataDownloadEmsLog:
		logType, err := resolveEmsLogType(item.Subtype)
		if err != nil {
			return err
		}
		data, filename, err = fetchEmsLog(ctx, item, logType)
		if err != nil {
			return err
		}
		contentType = "application/octet-stream"
	default:
		return fmt.Errorf("%s is not a binary document kind", item.Kind)
	}
	contentLock, err := cache.Lock(ctx, "data-download:content:"+documentID, 10*time.Minute, 200*time.Millisecond, 30*time.Second)
	if err != nil {
		return err
	}
	defer contentLock.SafeRelease(context.Background())
	if err := lease.validate(ctx); err != nil {
		return err
	}
	objectKey := ""
	if len(data) > 0 {
		objectKey, err = minioRepo.PutBinaryDocument(ctx, documentID, filename, bytes.NewReader(data), int64(len(data)), contentType)
		if err != nil {
			return err
		}
		zap.L().Debug("[DataDownload] 二进制对象已上传 MinIO",
			zap.String("documentId", documentID), zap.String("kind", string(item.Kind)),
			zap.String("objectKey", objectKey), zap.Int("bytes", len(data)))
	}
	generation := int64(1)
	createdAt := fetchedAt
	oldObjectKey := ""
	if existing, err := mongoRepo.GetDataDownloadDocument(ctx, documentID); err == nil {
		generation = existing.ContentGeneration + 1
		createdAt = existing.CreatedAt
		oldObjectKey = existing.BinaryObjectKey
	} else if err != mongo.ErrNoDocuments {
		return err
	}
	sum := sha256.Sum256(data)
	state := domain.DataDownloadMutable
	var frozenAt *time.Time
	loc := dataDownloadLocation(item.Timezone)
	day, _ := time.ParseInLocation(time.DateOnly, item.Date, loc)
	if !time.Now().In(loc).Before(day.AddDate(0, 0, 2)) {
		now := time.Now().UTC()
		state, frozenAt = domain.DataDownloadFrozen, &now
	}
	document := &domain.DataDownloadDocument{ID: documentID, RequestKey: lease.requestKey, Kind: item.Kind, SiteSN: item.SiteSN, DeviceSN: item.DeviceSN, Date: item.Date, Subtype: item.Subtype, DeviceType: item.DeviceType, Timezone: item.Timezone, State: state, Empty: len(data) == 0, ContentGeneration: generation, BinaryObjectKey: objectKey, Checksum: hex.EncodeToString(sum[:]), SourceFetchedAt: fetchedAt, FrozenAt: frozenAt, CreatedAt: createdAt, UpdatedAt: fetchedAt}
	if err := mongoRepo.UpsertDataDownloadDocument(ctx, document); err != nil {
		return err
	}
	if oldObjectKey != "" && oldObjectKey != objectKey {
		_ = minioRepo.Delete(ctx, oldObjectKey)
	}
	if err := lease.complete(ctx); err != nil {
		return err
	}
	return linkDataDownloadItem(ctx, item, document, document.Empty)
}

// DataDownloadSNDeferred keeps another file for this SN in the queue without
// consuming an Asynq retry or another worker while the current call runs.
type DataDownloadSNDeferred struct{ After time.Duration }

func (e *DataDownloadSNDeferred) Error() string { return "EMS log call for SN is busy" }

func fetchEmsLog(ctx context.Context, item *domain.DataDownloadItem, logType int) ([]byte, string, error) {
	slot, ok, err := jobs.TryAcquireEmsLogCallSlot(ctx, item.SiteSN)
	if err != nil {
		return nil, "", err
	}
	if !ok {
		return nil, "", &DataDownloadSNDeferred{After: 15 * time.Second}
	}
	defer slot.Release()
	callCtx, cancel := context.WithTimeout(slot.Context(), jobs.EmsLogCallTimeout)
	defer cancel()
	baselineID, err := repositories.LatestEmsLogCallRecordIDContext(callCtx, item.SiteSN, item.Date, logType)
	if callCtx.Err() != nil {
		if errors.Is(callCtx.Err(), context.DeadlineExceeded) {
			return nil, "", jobs.ErrEmsLogCallTimeout
		}
		return nil, "", jobs.ErrEmsLogCallSlotLost
	}
	if err != nil {
		return nil, "", err
	}
	err = repositories.QueryCallEmsLogContext(callCtx, item.SiteSN, item.Date, logType)
	if errors.Is(callCtx.Err(), context.DeadlineExceeded) {
		return nil, "", jobs.ErrEmsLogCallTimeout
	}
	if callCtx.Err() != nil && ctx.Err() == nil {
		return nil, "", jobs.ErrEmsLogCallSlotLost
	}
	if err != nil {
		return nil, "", err
	}
	zap.L().Debug("[DataDownload] EmsLog 已召唤，等待生成",
		zap.String("sn", item.SiteSN), zap.String("date", item.Date), zap.Int("logType", logType))
	record, err := waitEmsLog(callCtx, item.SiteSN, item.Date, logType, baselineID)
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) {
			return nil, "", jobs.ErrEmsLogCallTimeout
		}
		if errors.Is(err, context.Canceled) && ctx.Err() == nil {
			return nil, "", jobs.ErrEmsLogCallSlotLost
		}
		return nil, "", err
	}
	data, err := repositories.QueryDownloadEmsLogFileContext(callCtx, record.FileUrl)
	if errors.Is(callCtx.Err(), context.DeadlineExceeded) {
		return nil, "", jobs.ErrEmsLogCallTimeout
	}
	if callCtx.Err() != nil && ctx.Err() == nil {
		return nil, "", jobs.ErrEmsLogCallSlotLost
	}
	if err != nil {
		return nil, "", err
	}
	zap.L().Debug("[DataDownload] EmsLog 文件已下载到内存",
		zap.String("sn", item.SiteSN), zap.String("date", item.Date),
		zap.Int("bytes", len(data)), zap.String("fileUrl", record.FileUrl))
	filename := path.Base(record.FileUrl)
	if filename == "." || filename == "" {
		filename = fmt.Sprintf("EmsLog_%s_%s.tar.gz", item.Date, item.Subtype)
	}
	return data, filename, nil
}

func waitEmsLog(ctx context.Context, sn, date string, logType int, baselineID int64) (*domain.EmsLogRecordInfo, error) {
	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-ticker.C:
			if err := ctx.Err(); err != nil {
				return nil, err
			}
			record, err := repositories.QueryEmsLogCallRecordAfterContext(ctx, sn, date, logType, baselineID)
			if ctx.Err() != nil {
				return nil, ctx.Err()
			}
			if err != nil {
				continue
			}
			if record.Status < 0 {
				return nil, fmt.Errorf("EMS log request failed")
			}
			if record.Status == 1 {
				return record, nil
			}
		}
	}
}

func stableDataDownloadDocumentID(requestKey string) string {
	sum := sha256.Sum256([]byte(requestKey))
	return hex.EncodeToString(sum[:])
}

func linkDataDownloadItem(ctx context.Context, item *domain.DataDownloadItem, document *domain.DataDownloadDocument, empty bool) error {
	task, err := mongoRepo.GetDataDownloadTask(ctx, item.TaskID)
	if err != nil {
		return err
	}
	if task.Status == domain.DataDownloadCancelling || task.Status == domain.DataDownloadCancelled {
		return nil
	}
	logicalPath := fmt.Sprintf("%s/%s/%s", item.SiteSN, item.Kind, item.Date)
	if item.Kind == domain.DataDownloadFaultRecord {
		logicalPath = fmt.Sprintf("%s/%s/%s-%s/%s", item.SiteSN, item.Kind, item.DeviceSN, item.DeviceType, item.Date)
	}
	return mongoRepo.LinkDataDownloadDocument(ctx, &domain.DataDownloadTaskDocument{TaskID: item.TaskID, ItemKey: item.ItemKey, DocumentID: document.ID, LogicalPath: logicalPath, Empty: empty, CreatedAt: time.Now().UTC()})
}

// dataDownloadEmsLogTypeMap 将 EMS 日志的字符串枚举映射为上游要求的数字 logType，
// 与 BatchDownload(fwhTaskCallFile.go)的 mapEmsLogType 保持一致。
var dataDownloadEmsLogTypeMap = map[string]int{
	"EmsRunLog":     4,  // IBG运行日志
	"EmsBattLog":    7,  // 电池日志
	"EmsMonitorLog": 9,  // 设备系统资源监控日志
	"Ems4gLog":      10, // 4G日志
	"EmsAwsLog":     11, // AWS日志
	"EmsMainLog":    12, // 主进程日志
}

// resolveEmsLogType 优先按字符串枚举映射为数字 logType；
// 若传入的本身就是数字字符串(如 "4")则兼容解析，保证历史数据与测试可用。
func resolveEmsLogType(subtype string) (int, error) {
	if n, ok := dataDownloadEmsLogTypeMap[subtype]; ok {
		return n, nil
	}
	if n, err := strconv.Atoi(subtype); err == nil {
		return n, nil
	}
	return 0, fmt.Errorf("invalid EMS log type %q", subtype)
}