FranklinWH-Toolbox/service/services/batchDownloadServices.go

package services

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"compress/gzip"
	"fmt"
	"fwh-toolbox-server/domain"
	"fwh-toolbox-server/dto"
	"fwh-toolbox-server/repositories"
	"fwh-toolbox-server/taskQueue"
	"fwh-toolbox-server/utils/minio"
	"io"
	"reflect"
	"strings"

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

// 从数据库获取批量数据下载列表
func GetBatchDownloadManageList(page, size uint, sn, progress *string) (*dto.ResPage[domain.BatchDownloadInfo], error) {
	return repositories.GetBatchDownloadManageList(page, size, sn, progress)
}

// 插入一条批量数据下载配置到数据库
func InsertBatchDownloadConfig(writeConfig domain.BatchDownloadConfigWrite) error {
	id, err := repositories.InsertBatchDownloadConfig(writeConfig)
	if err != nil {
		return err
	}

	taskQueue.EnqueueBatchDownloadTask(domain.TaskBatchDownload{
		Id:         id,
		Sn:         writeConfig.Sn,
		TimePeriod: writeConfig.TimePeriod,
		Options:    writeConfig.Options,
	})

	return nil
}

// 更新一条批量数据下载配置到数据库
func UpdateBatchDownloadDescription(id primitive.ObjectID, desc string) error {
	err := repositories.UpdateBatchDownloadDescription(id, desc)
	if err != nil {
		return err
	}

	return nil
}

// 设置一个批量数据的owner
func UpdateBatchDownloadOwner(id primitive.ObjectID, owner []string) error {
	err := repositories.UpdateBatchDownloadOwner(id, owner)
	if err != nil {
		return err
	}

	return nil
}

// 删除一条批量数据下载配置
func DeleteBatchDownloadConfig(id primitive.ObjectID) error {
	err := repositories.DeleteBatchDownloadConfig(id)
	if err != nil {
		return err
	}

	// TODO 需删除相关数据

	return nil
}

// 从数据库获取批量数据下载文件地址并打包成zip文件
func GetBatchDownloadData(id primitive.ObjectID, writer *bufio.Writer) error {
	dataPathMap, err := repositories.GetBatchDownloadData(id)
	if err != nil {
		return err
	}

	// Pipe：让 ZIP 写入端与 Fiber 输出端实时对接
	pr, pw := io.Pipe()

	go func() {
		defer pw.Close()

		fmt.Println("Start Zip.")
		zw := zip.NewWriter(pw)
		defer zw.Close()

		var InsertZipFile = func(path string) error {
			obj, err := minio.GetObject(path)
			if err != nil {
				zap.L().Error("获取对象失败", zap.String("objName", path), zap.String("err", err.Error()))
				return err
			}
			defer obj.Close()

			zipFileWriter, err := zw.Create(path)
			if err != nil {
				zap.L().Error("创建zip条目失败", zap.String("path", path), zap.String("err", err.Error()))
				return err
			}

			if _, err := io.Copy(zipFileWriter, obj); err != nil {
				zap.L().Error("写入zip失败", zap.String("path", path), zap.String("err", err.Error()))
				return err
			}
			return nil
		}

		// map[sn]batchData
		for _, batchData := range dataPathMap {
			// PowerFlow
			// map[date]path
			if batchData.PowerFlow != nil {
				for _, path := range *batchData.PowerFlow {
					fmt.Println("InsertZipFile: ", path)
					if err := InsertZipFile(path); err != nil {
						continue
					}
				}
			}

			// MqttLogs
			// map[date]logs
			if batchData.MqttLogs != nil {
				for _, logs := range *batchData.MqttLogs {
					// map[logType]path
					for _, path := range logs {
						fmt.Println("InsertZipFile: ", path)
						if err := InsertZipFile(path); err != nil {
							continue
						}
					}
				}
			}

			// EventAlarm
			// map[date]path
			if batchData.EventAlarm != nil {
				for _, path := range *batchData.EventAlarm {
					fmt.Println("InsertZipFile: ", path)
					if err := InsertZipFile(path); err != nil {
						continue
					}
				}
			}

			// EmsLogs
			// map[date]logs
			if batchData.EmsLogs != nil {
				for _, logs := range *batchData.EmsLogs {
					// map[logType]path
					for _, path := range logs {
						fmt.Println("InsertZipFile: ", path)
						if err := InsertZipFile(path); err != nil {
							continue
						}
					}
				}
			}

			// FaultRecord
			// map[date]path
			if batchData.FaultRecord != nil {
				for _, path := range *batchData.FaultRecord {
					fmt.Println("InsertZipFile: ", path)
					if err := InsertZipFile(path); err != nil {
						continue
					}
				}
			}

			// UpdateLog
			if batchData.UpdateLog != nil && *batchData.UpdateLog != "" {
				fmt.Println("InsertZipFile: ", batchData.UpdateLog)
				if err := InsertZipFile(*batchData.UpdateLog); err != nil {
					continue
				}
			}

			// Settings
			if batchData.Settings != nil && *batchData.Settings != "" {
				fmt.Println("InsertZipFile: ", batchData.Settings)
				if err := InsertZipFile(*batchData.Settings); err != nil {
					continue
				}
			}
		}
		fmt.Println("End Zip.")
	}()

	// Fiber 负责读取 pr 并发送给客户端（消费者）
	_, err = io.Copy(writer, pr) // 等 pipe EOF 自动结束
	if err != nil {
		return err
	}

	return nil
}

func createBatchDownloadTree(data map[string]domain.BatchDownloadData) []dto.BatchDownloadStructNode {
	refVal := reflect.ValueOf(data)

	node := []dto.BatchDownloadStructNode{}
	if refVal.Kind() != reflect.Map {
		return node
	}

	// map 一定是目录
	for _, key := range refVal.MapKeys() {
		val := refVal.MapIndex(key)

		// 👉 value 是 BatchDownloadData Struct
		if val.Kind() != reflect.Struct {
			continue
		}

		child := createBatchDownloadTreeNode(key.String(), val)
		if child != nil {
			node = append(node, *child)
		}
	}

	return node
}

func createBatchDownloadTreeNode(name string, v reflect.Value) *dto.BatchDownloadStructNode {
	if !v.IsValid() {
		return nil
	}

	// 解指针
	for v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return nil
		}
		v = v.Elem()
	}

	node := &dto.BatchDownloadStructNode{
		Name: name,
	}

	switch v.Kind() {

	case reflect.Struct:
		// struct 一定是目录
		for i := 0; i < v.NumField(); i++ {
			field := v.Type().Field(i)
			val := v.Field(i)

			child := createBatchDownloadTreeNode(field.Name, val)
			if child != nil {
				node.Children = append(node.Children, child)
			}
		}

	case reflect.Map:
		// map 一定是目录
		for _, key := range v.MapKeys() {
			val := v.MapIndex(key)

			// 👉 value 是 string：文件
			if val.Kind() == reflect.String {
				node.Children = append(node.Children, &dto.BatchDownloadStructNode{
					Name:   key.String(),
					IsFile: false,
					Children: []*dto.BatchDownloadStructNode{
						{
							Name:   val.String(),
							IsFile: true,
						},
					},
				})
				continue
			}

			// 👉 value 是 map[string]string：下一层 value 才是文件
			child := createBatchDownloadTreeNode(key.String(), val)
			if child != nil {
				node.Children = append(node.Children, child)
			}
		}

	case reflect.String:
		// *string 或 string ⇒ 文件
		node.IsFile = true
		return node
	}

	// 目录节点
	node.IsFile = false

	// 空目录不展示
	if len(node.Children) == 0 {
		return nil
	}

	return node
}

// 获取指定ID的文件目录结构
func GetBatchDownloadPath(id primitive.ObjectID) ([]dto.BatchDownloadStructNode, error) {
	cfg, err := repositories.GetBatchDownloadData(id)
	if err != nil {
		return nil, err
	}

	tree := createBatchDownloadTree(cfg)
	return tree, nil
}

func createTarTree(tr *tar.Reader) (*dto.BatchDownloadTarNode, error) {
	root := &dto.BatchDownloadTarNode{Name: "/", IsDir: true}

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		parts := strings.Split(strings.Trim(hdr.Name, "/"), "/")
		current := root

		for i, p := range parts {
			isLast := i == len(parts)-1
			var child *dto.BatchDownloadTarNode

			for _, c := range current.Children {
				if c.Name == p {
					child = c
					break
				}
			}

			if child == nil {
				child = &dto.BatchDownloadTarNode{
					Name:  p,
					Path:  strings.Join(parts[:i+1], "/"),
					IsDir: !isLast || hdr.Typeflag == tar.TypeDir,
				}
				current.Children = append(current.Children, child)
			}

			current = child
		}
	}

	return root, nil
}

// 获取指定文件的目录结构
func GetBatchDownloadFileTree(filePath string) (*dto.BatchDownloadTarNode, error) {
	obj, err := minio.GetObject(filePath)
	if err != nil {
		return nil, err
	}

	gz, err := gzip.NewReader(obj)
	if err != nil {
		return nil, err
	}

	tr := tar.NewReader(gz)

	return createTarTree(tr)
}

// 下载指定文件
func GetBatchDownloadFileDownload(filePath, innerPath string, writer *bufio.Writer) error {
	obj, err := minio.GetObject(filePath)
	if err != nil {
		zap.L().Error("Minio get object filed.", zap.String("filePath", filePath), zap.String("err", err.Error()))
		return err
	}
	defer obj.Close()

	if strings.HasSuffix(filePath, ".tar.gz") && innerPath != "" {
		// 下载.tar.gz包中的文件
		gz, err := gzip.NewReader(obj)
		if err != nil {
			zap.L().Error("gz file read filed.", zap.String("filePath", filePath), zap.String("err", err.Error()))
			return err
		}
		defer gz.Close()

		tr := tar.NewReader(gz)

		for {
			hdr, err := tr.Next()
			if err == io.EOF {
				break
			}
			if err != nil {
				return err
			}

			if hdr.Typeflag != tar.TypeReg {
				continue
			}

			if hdr.Name != innerPath {
				continue
			}

			// ⭐ 核心：流式拷贝
			_, err = io.Copy(writer, tr)
			return err
		}

		zap.L().Error("inner file not found in tar.", zap.String("filePath", filePath), zap.String("innerPath", innerPath), zap.String("err", err.Error()))
		return fmt.Errorf("inner file not found in tar: %s", innerPath)
	}

	// ⭐ 核心：流式转发
	_, err = io.Copy(writer, obj)
	if err != nil {
		// 客户端中断时，这里很常见
		zap.L().Error("Minio get object download interrupted.", zap.String("filePath", filePath), zap.String("err", err.Error()))
		return err
	}

	return nil
}