FranklinWH-Toolbox/service/task/queue/fwhQueue/fwhTaskBatchDownload.go

package fwhQueue

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"fwh-toolbox-server/cache"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/dto"
	"fwh-toolbox-server/jobs"
	"fwh-toolbox-server/repositories"
	"strconv"
	"strings"
	"time"

	"github.com/hibiken/asynq"
	"go.uber.org/zap"
)

// handleBatchDownload 是 BatchDownload 主任务的 Asynq handler
func handleBatchDownload(ctx context.Context, task *asynq.Task) error {
	var payload domain.TaskBatchDownload
	if err := json.Unmarshal(task.Payload(), &payload); err != nil {
		return fmt.Errorf("unmarshal BatchDownload payload: %w", err)
	}

	mainTaskId := payload.Id.Hex()

	// 更新运行状态
	if err := repositories.UpdateBatchDownloadResult(payload.Id, domain.ProgressRunning, []string{}); err != nil {
		zap.L().Error("[BatchDownload] 更新状态失败(Running)", zap.String("id", mainTaskId))
		return err
	}

	// 执行主逻辑
	reasons, err := executeBatchDownload(ctx, payload)

	// 记录结果
	if err != nil {
		zap.L().Error("[BatchDownload] 任务执行失败", zap.String("id", mainTaskId), zap.String("err", err.Error()))
		repositories.UpdateBatchDownloadResult(payload.Id, domain.ProgressError, []string{err.Error()})
		return err
	}

	if len(reasons) == 0 {
		zap.L().Info("[BatchDownload] 任务执行成功", zap.String("id", mainTaskId))
		repositories.UpdateBatchDownloadResult(payload.Id, domain.ProgressComplete, []string{})
	} else {
		zap.L().Warn("[BatchDownload] 任务部分失败", zap.String("id", mainTaskId), zap.Strings("reasons", reasons))
		repositories.UpdateBatchDownloadResult(payload.Id, domain.ProgressFailed, reasons)
	}

	return nil // 返回 nil 让 Asynq 标记任务为 completed，不触发重试
}

// executeBatchDownload 执行批量下载逻辑
func executeBatchDownload(ctx context.Context, payload domain.TaskBatchDownload) ([]string, error) {
	mainTaskId := payload.Id.Hex()

	// 解析时间范围
	startTime, err := time.Parse(time.DateOnly, payload.TimePeriod[0])
	if err != nil {
		return nil, fmt.Errorf("startTime 格式错误: %w", err)
	}
	endTime, err := time.Parse(time.DateOnly, payload.TimePeriod[1])
	if err != nil {
		return nil, fmt.Errorf("endTime 格式错误: %w", err)
	}
	endTime = endTime.Add(23*time.Hour + 59*time.Minute + 59*time.Second)
	if startTime.After(endTime) {
		return nil, errors.New("startTime > endTime")
	}

	// 按天拆分时间段
	periods := splitTimePeriods(startTime, endTime)

	// ==================== 阶段1：创建文件下载子任务 ====================
	callFileTasks := buildCallFileTasks(payload, periods, startTime, endTime)
	totalSubTasks := len(callFileTasks)

	if totalSubTasks > 0 {
		// 初始化 Redis 计数器（待完成数）
		counterKey := callFileCounterPrefix + mainTaskId
		failuresKey := callFileFailuresPrefix + mainTaskId

		cache.Set(ctx, counterKey, totalSubTasks, 2*time.Hour)
		cache.Delete(ctx, failuresKey)

		// 入队所有子任务
		for i, cfTask := range callFileTasks {
			taskPayload, _ := json.Marshal(cfTask)
			subTaskID := fmt.Sprintf("callfile-%s-%d", mainTaskId, i)
			t := asynq.NewTask(TaskTypeCallFile, taskPayload,
				asynq.MaxRetry(3),
				asynq.Timeout(35*time.Minute),
				asynq.Queue(QueueDownload),
				asynq.TaskID(subTaskID),
			)
			if _, err := TaskClient.Enqueue(t); err != nil {
				zap.L().Error("[BatchDownload] 子任务入队失败",
					zap.String("id", mainTaskId),
					zap.Int("index", i),
					zap.String("err", err.Error()),
				)
			}
		}
	}

	// ==================== 阶段2：同步处理其他数据任务 ====================
	if err := executeSyncTasks(ctx, payload, periods); err != nil {
		return nil, err
	}

	// ==================== 阶段3：等待子任务完成 ====================
	if totalSubTasks > 0 {
		failReasons, err := waitForCallFileTasks(ctx, mainTaskId)
		if err != nil {
			return nil, err
		}
		return failReasons, nil
	}

	return nil, nil
}

// callFilePayload 是传递给子任务的 payload 结构
type callFilePayload struct {
	MainTaskId string                           `json:"mainTaskId"`
	TaskInfo   domain.BatchDownloadFileTaskInfo `json:"taskInfo"`
}

// buildCallFileTasks 构建所有需要异步下载的文件任务列表
func buildCallFileTasks(payload domain.TaskBatchDownload, periods []timePeriod, startTime, endTime time.Time) []callFilePayload {
	mainTaskId := payload.Id.Hex()
	var tasks []callFilePayload

	for _, sn := range payload.Sn {
		// EmsLog 任务
		for _, period := range periods {
			dateStr := period.Start.Format(time.DateOnly)
			for _, logType := range payload.Options.EmsLogs {
				tasks = append(tasks, callFilePayload{
					MainTaskId: mainTaskId,
					TaskInfo: domain.BatchDownloadFileTaskInfo{
						DevType:     -1,
						Sn:          sn,
						Date:        dateStr,
						FileType:    jobs.FileTypeEmsLog,
						FileSupType: logType,
					},
				})
			}
		}

		// FaultRecord 任务
		if payload.Options.FaultRecord {
			frTasks := buildFaultRecordTasks(mainTaskId, sn, startTime, endTime)
			tasks = append(tasks, frTasks...)
		}
	}

	return tasks
}

// buildFaultRecordTasks 为单个 SN 构建故障录波下载任务
func buildFaultRecordTasks(mainTaskId string, sn string, startTime, endTime time.Time) []callFilePayload {
	var tasks []callFilePayload

	// 查询 aGate 可召唤列表
	aGateTasks := buildFaultRecordTasksForDev(mainTaskId, 1, sn, startTime, endTime)
	tasks = append(tasks, aGateTasks...)

	// 查询 aPower 子设备
	siteBaseInfoList, err := repositories.QuerySiteBaseInfoList(dto.SiteBaseInfoPara{
		SiteList: []string{sn},
		PageSize: 10,
		PageNum:  1,
	})
	if err != nil {
		zap.L().Error("[BatchDownload] 查询站点基本信息失败", zap.String("sn", sn), zap.String("err", err.Error()))
		return tasks
	}

	if siteBaseInfoList != nil && len(*siteBaseInfoList.Result) > 0 {
		aPowerSnList := strings.Split((*siteBaseInfoList.Result)[0].FhpSn, ",")
		for _, aPowerSn := range aPowerSnList {
			aPowerSn = strings.TrimSpace(aPowerSn)
			if aPowerSn == "" {
				continue
			}
			aPowerTasks := buildFaultRecordTasksForDev(mainTaskId, 0, aPowerSn, startTime, endTime)
			tasks = append(tasks, aPowerTasks...)
		}
	}

	return tasks
}

// buildFaultRecordTasksForDev 为指定设备构建故障录波下载任务
func buildFaultRecordTasksForDev(mainTaskId string, devType int, sn string, startTime, endTime time.Time) []callFilePayload {
	callInfoList, err := repositories.QueryFaultRecordCallableList(devType, sn)
	if err != nil {
		zap.L().Error("[BatchDownload] 查询故障录波召唤选项失败", zap.String("sn", sn), zap.String("err", err.Error()))
		return nil
	}
	if len(callInfoList) == 0 {
		return nil
	}

	// 筛选日期范围内的记录
	var indexList string
	for _, info := range callInfoList {
		failureTime, err := parseFaultRecordTime(info.FailureTime)
		if err != nil {
			continue
		}
		inRange := (failureTime.Equal(startTime) || failureTime.After(startTime)) &&
			(failureTime.Equal(endTime) || failureTime.Before(endTime))
		if inRange {
			if indexList == "" {
				indexList = strconv.Itoa(info.Idx)
			} else {
				indexList = indexList + "," + strconv.Itoa(info.Idx)
			}
		}
	}

	// 没有符合日期范围的记录，不创建任务
	if indexList == "" {
		return nil
	}

	return []callFilePayload{{
		MainTaskId: mainTaskId,
		TaskInfo: domain.BatchDownloadFileTaskInfo{
			DevType:     devType,
			Sn:          sn,
			Date:        startTime.Format(time.DateTime) + "," + endTime.Format(time.DateTime),
			FileType:    jobs.FileTypeFaultRecord,
			FileSupType: indexList,
		},
	}}
}

// executeSyncTasks 同步执行非文件下载类任务（PowerFlow、MqttLogs、EventAlarm 等）
func executeSyncTasks(ctx context.Context, payload domain.TaskBatchDownload, periods []timePeriod) error {
	for _, sn := range payload.Sn {
		// PowerFlow
		if payload.Options.PowerFlow {
			for _, period := range periods {
				dateStr := period.Start.Format(time.DateOnly)
				res, err := repositories.QuerySitePowerFlow(sn, period.Start)
				if err != nil {
					zap.L().Error("[BatchDownload] 查询功率流数据失败", zap.String("sn", sn), zap.String("date", dateStr), zap.String("err", err.Error()))
					return err
				}
				objName, err := jobs.SaveBatchDownloadPowerFlow(sn, dateStr, *res.Result)
				if err != nil {
					return err
				}
				if err := repositories.InsertBatchDownloadPathPowerFlow(payload.Id, sn, dateStr, objName); err != nil {
					return err
				}
			}
		}

		// MqttLogs
		for _, logType := range payload.Options.MqttLogs {
			for _, period := range periods {
				logList, err := repositories.QueryMqttLogListOnTimeRange(sn, domain.MqttCmdType(logType), period.Start, period.End)
				if err != nil {
					return err
				}
				dateStr := period.Start.Format(time.DateOnly)
				objName, err := jobs.SaveBatchDownloadMqttLog(sn, logType, dateStr, logList)
				if err != nil {
					return err
				}
				if err := repositories.InsertBatchDownloadPathMqttLogs(payload.Id, sn, logType, dateStr, objName); err != nil {
					return err
				}
			}
		}

		// EventAlarm
		if payload.Options.EventAlarm {
			for _, period := range periods {
				eventAlarmList, err := repositories.QueryEventOrAlarmListOnTimeRange(sn, period.Start, period.End, repositories.EventAndAlarm)
				if err != nil {
					return err
				}
				dateStr := period.Start.Format(time.DateOnly)
				objName, err := jobs.SaveBatchDownloadEventAlarm(sn, dateStr, eventAlarmList)
				if err != nil {
					return err
				}
				if err := repositories.InsertBatchDownloadPathEventAlarm(payload.Id, sn, dateStr, objName); err != nil {
					return err
				}
			}
		}

		// AlarmSnapshot
		if payload.Options.AlarmSnapshot {
			for _, period := range periods {
				dateStr := period.Start.Format(time.DateOnly)
				alarmSnapshotList, err := jobs.GetAlarmSnapshotData(ctx, sn, dateStr)
				if err != nil {
					return err
				}
				objName, err := jobs.SaveBatchDownloadAlarmSnapshot(sn, dateStr, alarmSnapshotList)
				if err != nil {
					return err
				}
				if err := repositories.InsertBatchDownloadPathAlarmSnapshot(payload.Id, sn, dateStr, objName); err != nil {
					return err
				}
			}
		}

		// UpdateLog
		if payload.Options.UpdateLog {
			if err := executeSyncUpdateLog(sn, payload); err != nil {
				return err
			}
		}
	}
	return nil
}

// waitForCallFileTasks 等待所有子任务完成，收集失败原因
func waitForCallFileTasks(ctx context.Context, mainTaskId string) ([]string, error) {
	counterKey := callFileCounterPrefix + mainTaskId
	failuresKey := callFileFailuresPrefix + mainTaskId

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	timeout := time.After(90 * time.Minute) // 90 分钟总超时

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-timeout:
			return nil, fmt.Errorf("等待子任务完成超时（90分钟）")
		case <-ticker.C:
			val, err := cache.Get(ctx, counterKey)
			if err != nil {
				// key 不存在或读取失败
				if err.Error() == "redis: nil" {
					// key 不存在，所有子任务已完成
					val = "0"
				} else {
					zap.L().Error("[BatchDownload] 读取计数器失败", zap.String("key", counterKey), zap.String("err", err.Error()))
					continue
				}
			}

			remaining, _ := strconv.Atoi(val)
			if remaining <= 0 {
				// 收集失败原因
				failures, err := cache.LRange(ctx, failuresKey, 0, -1)
				if err != nil {
					return nil, fmt.Errorf("读取失败原因列表失败: %w", err)
				}

				// 清理
				cache.Delete(ctx, counterKey, failuresKey)

				return failures, nil
			}
		}
	}
}

// ==================== 辅助函数 ====================

type timePeriod struct {
	Start time.Time
	End   time.Time
}

func splitTimePeriods(startTime, endTime time.Time) []timePeriod {
	var periods []timePeriod
	curr := startTime
	for curr.Before(endTime) {
		dayEnd := time.Date(curr.Year(), curr.Month(), curr.Day(), 23, 59, 59, 0, curr.Location())
		if dayEnd.After(endTime) {
			dayEnd = endTime
		}
		periods = append(periods, timePeriod{Start: curr, End: dayEnd})
		curr = time.Date(curr.Year(), curr.Month(), curr.Day()+1, 0, 0, 0, 0, curr.Location())
	}
	return periods
}

func parseFaultRecordTime(failureTime string) (time.Time, error) {
	parts := strings.Split(failureTime, " ")
	if len(parts) < 2 {
		return time.Time{}, fmt.Errorf("invalid failureTime format: %s", failureTime)
	}
	date := parts[0]
	timeStr := parts[1]

	d := strings.Split(date, "-")
	t := strings.Split(timeStr, ":")
	if len(d) < 3 || len(t) < 3 {
		return time.Time{}, fmt.Errorf("invalid failureTime format: %s", failureTime)
	}

	standardTime := fmt.Sprintf("%04s-%02s-%02s %02s:%02s:%02s", d[0], d[1], d[2], t[0], t[1], t[2])
	return time.Parse(time.DateTime, standardTime)
}

// executeSyncUpdateLog 同步执行升级日志任务（从原代码提取）
func executeSyncUpdateLog(sn string, payload domain.TaskBatchDownload) error {
	updateLogList := []domain.BatchDownloadSystemUpdateLogInfo{}

	systemUpdateLogList, err := repositories.QuerySystemUpdateLog(sn)
	if err != nil {
		return err
	}

	for _, logList := range systemUpdateLogList {
		var updateLogInfo domain.BatchDownloadSystemUpdateLogInfo
		updateLogInfo.OrderNo = logList.OrderNo
		updateLogInfo.GateSn = sn
		updateLogInfo.FileName = logList.Name

		upgradeTime, err := time.ParseInLocation(time.DateTime, logList.UpgradeTime, time.UTC)
		if err != nil {
			zap.L().Warn("[BatchDownload] 升级时间解析失败", zap.String("time", logList.UpgradeTime))
			continue
		}
		beijing, _ := time.LoadLocation("Asia/Shanghai")
		updateLogInfo.UpdateTime = upgradeTime.In(beijing).Format(time.DateTime)

		if strings.HasPrefix(logList.OrderNo, "CH") {
			updateLogInfo.TriggerType = "Auto"
		} else if strings.HasPrefix(logList.OrderNo, "DH") {
			updateLogInfo.TriggerType = "Manual"
		} else if strings.HasPrefix(logList.OrderNo, "SH") {
			updateLogInfo.TriggerType = "Span Auto"
		} else {
			updateLogInfo.TriggerType = "unknown"
		}

		switch logList.Result {
		case 1:
			updateLogInfo.Result = "Updating"
		case 2:
			updateLogInfo.Result = "Success"
		case 3:
			updateLogInfo.Result = "Fault"
		default:
			updateLogInfo.Result = fmt.Sprintf("unknown(%d)", logList.Result)
		}

		systemUpdateLogListDetail, err := repositories.QuerySystemUpdateLogDetail(logList.Id)
		if err != nil {
			return err
		}

		for _, logDetail := range systemUpdateLogListDetail {
			if logDetail.GatewayId == sn {
				for _, ibgVer := range logDetail.IbgVersionNew {
					if ibgVer.Children != nil {
						for _, childVerInfo := range ibgVer.Children {
							updateLogInfo.GateVerList = append(updateLogInfo.GateVerList, domain.BatchDownloadUpdateLogVersionInfo{
								Name: childVerInfo.Type, OldVer: childVerInfo.OldVersion, NewVer: childVerInfo.Version,
							})
						}
					} else {
						updateLogInfo.GateVerList = append(updateLogInfo.GateVerList, domain.BatchDownloadUpdateLogVersionInfo{
							Name: ibgVer.Type, OldVer: ibgVer.OldVersion, NewVer: ibgVer.Version,
						})
					}
				}
				for _, subInfo := range logDetail.SonNewDevices {
					var subDevInfo domain.BatchDownloadUpdateLogSubDevInfo
					subDevInfo.DevSn = subInfo.ApowerSn
					for _, firstItem := range subInfo.FirstList {
						for _, secondItem := range firstItem.Secondlist {
							for _, thirdItem := range secondItem.Thirdlist {
								subDevInfo.DevVerList = append(subDevInfo.DevVerList, domain.BatchDownloadUpdateLogVersionInfo{
									Name: thirdItem.SubType, OldVer: thirdItem.OldVersion, NewVer: thirdItem.Version,
								})
							}
						}
					}
					updateLogInfo.SubDevList = append(updateLogInfo.SubDevList, subDevInfo)
				}
				break
			}
		}
		updateLogList = append(updateLogList, updateLogInfo)
	}

	// aHub 升级日志
	aHubUpdateLogList, err := repositories.QueryAHubUpdateLog(sn)
	if err != nil {
		return err
	}

	var aHubUpdateLogData []domain.BatchDownloadAHubUpdateLogInfo
	for _, aHubVerInfo := range aHubUpdateLogList {
		var info domain.BatchDownloadAHubUpdateLogInfo
		info.OrderNo = aHubVerInfo.OrderNo
		info.GateSn = sn
		info.AHubSn = aHubVerInfo.AHubSn

		if strings.HasPrefix(aHubVerInfo.OrderNo, "CH") {
			info.TriggerType = "Auto"
		} else if strings.HasPrefix(aHubVerInfo.OrderNo, "DH") {
			info.TriggerType = "Manual"
		} else if strings.HasPrefix(aHubVerInfo.OrderNo, "SH") {
			info.TriggerType = "Span Auto"
		} else {
			info.TriggerType = "unknown"
		}

		switch aHubVerInfo.UpgradeResult {
		case 0:
			info.Result = "Success"
		case -1:
			info.Result = "Fault"
		default:
			info.Result = fmt.Sprintf("unknown(%d)", aHubVerInfo.UpgradeResult)
		}

		for _, detailInfo := range aHubVerInfo.DetailList {
			info.DetailList = append(info.DetailList, domain.BatchDownloadUpdateLogAHubVersionInfo{
				FileName:   detailInfo.UpgradeFileName,
				FileOrder:  detailInfo.UpgradeFileOrder,
				Result:     detailInfo.UpgradeResult,
				ResultDesc: detailInfo.UpgradeResultDesc,
				OldVer:     detailInfo.OldVersion,
				NewVer:     detailInfo.NewVersion,
			})
		}
		aHubUpdateLogData = append(aHubUpdateLogData, info)
	}

	objName, err := jobs.SaveBatchDownloadUpdateLog(sn, updateLogList, aHubUpdateLogData)
	if err != nil {
		return err
	}
	return repositories.InsertBatchDownloadPathUpdateLog(payload.Id, sn, objName)
}