FranklinWH-Toolbox/service/repositories/mongo/batchDownloadRepo.go

package mongo

import (
	"context"
	"fwh-toolbox-server/database"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/dto"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const (
	CollectionName_BatchDownloadConfig = "BatchDownloadConfig"
)

// 从数据库获取批量下载数据列表
func GetBatchDownloadManageList(page, size uint, sn, progress *string) (*dto.ResPage[domain.BatchDownloadInfo], error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	filter := bson.M{}
	// 条件
	if sn != nil {
		filter["sn"] = *sn
	}
	if progress != nil {
		filter["progress"] = *progress
	}

	findOptions := options.Find().
		SetProjection(bson.M{
			"_id":         1,
			"sn":          1,
			"description": 1,
			"timePeriod":  1,
			"options":     1,
			"progress":    1,
			"failReasons": 1,
			"items":       1,
			"totalCount":  1,
			"doneCount":   1,
		}).
		SetSkip(int64((page - 1) * size)).
		SetLimit(int64(size)).
		SetSort(bson.D{
			{Key: "_id", Value: -1}, // 倒序
		})

	cursor, err := collection.Find(ctx, filter, findOptions)
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var configList []domain.BatchDownloadInfo
	for cursor.Next(ctx) {
		var cfg domain.BatchDownloadInfo
		if err := cursor.Decode(&cfg); err != nil {
			return nil, err
		}
		configList = append(configList, cfg)
	}

	// 查询总数
	total, err := collection.CountDocuments(ctx, filter)
	if err != nil {
		return nil, err
	}

	res := dto.ResPage[domain.BatchDownloadInfo]{
		Data:     configList,
		PageSize: size,
		PageNum:  page,
		Total:    uint(total),
	}

	return &res, nil
}

// 从数据库获取一条批量下载数据配置
func GetBatchDownloadConfig(id primitive.ObjectID) (*domain.BatchDownloadConfig, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)
	var cfg domain.BatchDownloadConfig
	err := collection.FindOne(ctx, bson.M{"_id": id}).Decode(&cfg)
	if err != nil {
		return nil, err
	}

	return &cfg, nil
}

// 插入一条批量下载配置到数据库
func InsertBatchDownloadConfig(writeConfig domain.BatchDownloadConfigWrite) (primitive.ObjectID, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	config := domain.BatchDownloadConfig{
		Sn:          writeConfig.Sn,
		Description: writeConfig.Description,
		TimePeriod:  writeConfig.TimePeriod,
		Options:     writeConfig.Options,
		Progress:    domain.ProgressPending,
		Data:        map[string]domain.BatchDownloadData{},
		FailReasons: []string{},
	}

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	res, err := collection.InsertOne(ctx, config)
	if err != nil {
		return primitive.ObjectID{}, err
	}

	return res.InsertedID.(primitive.ObjectID), nil
}

// 更新一条批量下载描述描述到数据库
func UpdateBatchDownloadDescription(id primitive.ObjectID, desc string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新全部字段
		"$set": bson.M{
			"description": desc,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 更新一条批量下载数据owner到数据库
func UpdateBatchDownloadOwner(id primitive.ObjectID, owner []string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新全部字段
		"$set": bson.M{
			"owner": owner,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 更新一条批量下载结果到数据库
func UpdateBatchDownloadResult(id primitive.ObjectID, result domain.Progress, reasons []string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新非progress字段
		"$set": bson.M{
			"progress":    result,
			"failReasons": reasons,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 删除一条批量下载配置
func DeleteBatchDownloadConfig(id primitive.ObjectID) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	filter := bson.M{"_id": id} // 删除条件
	collection := database.GetCollection(CollectionName_BatchDownloadConfig)
	_, err := collection.DeleteOne(ctx, filter)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据
func InsertBatchDownloadData(id primitive.ObjectID, data domain.BatchDownloadData) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data": data,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 获取下载数据包路径
func GetBatchDownloadData(id primitive.ObjectID) (map[string]domain.BatchDownloadData, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)
	var cfg domain.BatchDownloadConfig
	err := collection.FindOne(ctx, bson.M{"_id": id}).Decode(&cfg)
	if err != nil {
		return nil, err
	}

	return cfg.Data, nil
}

// 插入批量下载结果数据 - PowerFlow
func InsertBatchDownloadPathPowerFlow(id primitive.ObjectID, sn string, date string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".powerFlow." + date: path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - MqttLogs
func InsertBatchDownloadPathMqttLogs(id primitive.ObjectID, sn string, date string, logType string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".mqttLogs." + date + "." + logType: path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - EventAlarm
func InsertBatchDownloadPathEventAlarm(id primitive.ObjectID, sn string, date string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".eventAlarm." + date: path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - AlarmSnapshot
func InsertBatchDownloadPathAlarmSnapshot(id primitive.ObjectID, sn string, date string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".alarmSnapshot." + date: path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - EmsLogs
func InsertBatchDownloadPathEmsLogs(id primitive.ObjectID, sn string, date string, logType string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".emsLogs." + date + "." + logType: path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - FaultRecord
func InsertBatchDownloadPathFaultRecord(id primitive.ObjectID, sn string, fileName string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".faultRecord." + fileName: path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - UpdateLog
func InsertBatchDownloadPathUpdateLog(id primitive.ObjectID, sn string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".updateLog": path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// 插入批量下载结果数据 - Settings
func InsertBatchDownloadPathSettings(id primitive.ObjectID, sn string, path string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	// _, err := collection.InsertOne(ctx, config)
	filter := bson.M{"_id": id} // 按ID匹配
	update := bson.M{
		// 如果存在就更新data字段
		"$set": bson.M{
			"data." + sn + ".settings": path,
		},
	}

	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}

	return nil
}

// ==================== 结构化 item 结果（B2 完成判定） ====================

// InitBatchDownloadItems 初始化 items 列表，并设置 totalCount / doneCount。
// doneCount 初始化为已经处于终态的 item 数量（构建期就 failed/skipped 的），
// 使得后续所有 item 转终态后 doneCount 恰好等于 totalCount。
func InitBatchDownloadItems(id primitive.ObjectID, items []domain.BatchDownloadItem) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	doneInit := 0
	now := time.Now()
	for i := range items {
		if items[i].UpdatedAt.IsZero() {
			items[i].UpdatedAt = now
		}
		if items[i].Status != domain.ItemPending {
			doneInit++
		}
	}

	filter := bson.M{"_id": id}
	update := bson.M{
		"$set": bson.M{
			"items":      items,
			"totalCount": len(items),
			"doneCount":  doneInit,
			"progress":   domain.ProgressRunning,
		},
	}
	opts := options.Update().SetUpsert(true)

	_, err := collection.UpdateOne(ctx, filter, update, opts)
	if err != nil {
		return err
	}
	return nil
}

// FinishBatchDownloadItem 原子地把某个 pending item 置为终态并递增 doneCount。
// 幂等：filter 要求该 item 当前为 pending，重复调用（Asynq 重试/重复回调）不会重复计数。
// 返回 (doneCount, totalCount, error)；若该 item 已非 pending（幂等跳过），返回 (-1, -1, nil)。
func FinishBatchDownloadItem(id primitive.ObjectID, key string, status domain.ItemStatus, stage, reason string) (int, int, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	filter := bson.M{
		"_id": id,
		"items": bson.M{
			"$elemMatch": bson.M{"key": key, "status": domain.ItemPending},
		},
	}
	update := bson.M{
		"$set": bson.M{
			"items.$.status":    status,
			"items.$.stage":     stage,
			"items.$.reason":    reason,
			"items.$.updatedAt": time.Now(),
		},
		"$inc": bson.M{"doneCount": 1},
	}
	opts := options.FindOneAndUpdate().SetReturnDocument(options.After)

	var doc domain.BatchDownloadConfig
	err := collection.FindOneAndUpdate(ctx, filter, update, opts).Decode(&doc)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			// item 已非 pending（重复回调/重试），幂等跳过
			return -1, -1, nil
		}
		return 0, 0, err
	}
	return doc.DoneCount, doc.TotalCount, nil
}

// GetBatchDownloadItems 读取某任务的全部 item 结果，用于聚合派生终态。
func GetBatchDownloadItems(id primitive.ObjectID) ([]domain.BatchDownloadItem, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	collection := database.GetCollection(CollectionName_BatchDownloadConfig)

	var cfg domain.BatchDownloadConfig
	err := collection.FindOne(ctx, bson.M{"_id": id},
		options.FindOne().SetProjection(bson.M{"items": 1})).Decode(&cfg)
	if err != nil {
		return nil, err
	}
	return cfg.Items, nil
}