pji_desktop_http/biz/handler/simulation_service/simulation_service.go

package simulation_service

import (
	"archive/zip"
	"context"
	"encoding/json"
	"fmt"
	"github.com/cloudwego/hertz/pkg/app"
	"github.com/cloudwego/hertz/pkg/protocol/consts"
	uuid "github.com/satori/go.uuid"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"pji_desktop_http/biz/dal/mysql"
	"pji_desktop_http/biz/model"
	"pji_desktop_http/common/config"
	"pji_desktop_http/common/config/c_log"
	"pji_desktop_http/common/entity"
	"pji_desktop_http/common/util"
	"strconv"
	"strings"
)

// DownloadSimulationZipFile 根据请求id从oss拉取并打包仿真测试所需要的文件
// @router /simulation/download/zipfile   [GET]
func DownloadSimulationZipFile(ctx context.Context, c *app.RequestContext) {
	id := c.Query("id")
	fmt.Println(id)

	// 根据id生成用于地图更新的压缩包
	filePath, tmpDir, err := generateSimulationZipById(ctx, id)
	if err != nil {
		c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
	}
	fmt.Println("filePath", filePath)

	// 检查文件是否存在
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		c.JSON(http.StatusNotFound, map[string]string{"error": "File not found"})
		return
	}

	// 打开文件
	f, err := os.Open(filePath)
	if err != nil {
		c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to open file"})
		return
	}
	defer f.Close()

	// 设置响应头
	c.Response.Header.Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filepath.Base(filePath)))
	c.Response.Header.Set("Content-Type", "binary/octet-stream")

	// 将文件流式传输回客户端
	data, err := io.ReadAll(f)
	if err != nil {
		panic(err)
	}
	if _, err := c.Write(data); err != nil {
		c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
		return
	}
	defer os.RemoveAll(tmpDir)
}

// 根据id生成用于仿真测试的压缩包
func generateSimulationZipById(ctx context.Context, id string) (file string, tmpDir string, err error) {
	// 根据id获取对应的oss文件列表
	allFileList, err := util.GetExactedMapFileById(id)

	//fmt.Println("Filtered Strings:", fileList)
	if err != nil {
		return
	}

	// 创建临时文件夹
	tmpDir, err = os.MkdirTemp("", "temp-download-*")

	fmt.Println("tmpDir:", tmpDir)
	if err != nil {
		fmt.Println("Error creating temporary directory:", err)
		return "", "", err
	}
	c_log.GlobalLogger.Info("创建下载-临时文件夹：", tmpDir)

	// 创建根文件夹（文件打包的根目录）
	baseDir := filepath.Join(tmpDir, "data")
	if err := os.Mkdir(baseDir, 0755); err != nil {
		fmt.Println("Error creating subdirectory:", err)
		return "", "", err
	}
	c_log.GlobalLogger.Info("创建文件打包根目录：", baseDir)

	// 从oss下载data.zip, map.pgm, map.yaml文件到根目录
	// 过滤特定后缀的文件列表
	simulationFileList := util.FilterBySuffixes(allFileList, config.SimulationFiltersuffixes...)
	fmt.Println("simulationFileList", simulationFileList)
	// 从oss下载文件到 根目录
	for _, file := range simulationFileList {
		err = config.OssBucket.GetObjectToFile(file, filepath.Join(baseDir, filepath.Base(file)))
		if err != nil {
			fmt.Println("Error downloading file - data.zip, map.pgm, map.yaml, map.bag:", err)
			return "", "", err
		}
	}
	c_log.GlobalLogger.Info("下载data.zip, map.pgm, map.yaml, map.bag文件到根目录 - 成功")

	// 下载world文件&map.stl文件
	// 查询状态
	world, err := mysql.QueryWorld(ctx, id)
	if err != nil || world == nil { // 记录不存在
		c_log.GlobalLogger.Info("world记录不存在 - 跳过")
	} else { // 记录存在
		// 下载world文件
		worldURL := world.WorldURL
		err = config.OssBucket.GetObjectToFile(worldURL, filepath.Join(baseDir, filepath.Base(worldURL)))
		if err != nil {
			fmt.Println("Error downloading world file:", err)
			return "", "", err
		}
		c_log.GlobalLogger.Info("下载world文件到根目录 - 成功")

		// 下载map.stl文件
		stlURL := world.StlURL
		if stlURL != nil && *stlURL != "" {
			err = config.OssBucket.GetObjectToFile(*stlURL, filepath.Join(baseDir, filepath.Base(*stlURL)))
			if err != nil {
				fmt.Println("Error downloading stl file:", err)
				return "", "", err
			}
			c_log.GlobalLogger.Info("下载map.stl文件到根目录 - 成功")
		} else {
			c_log.GlobalLogger.Info("map.stl文件不存在 - 跳过")
		}
	}

	// 下载原始bag
	rosField, err := util.GetRosFileById(id)
	if err != nil {
		fmt.Println("Error querying origin map's bag file:", err)
		return "", "", err
	}
	err = config.OssBucket.GetObjectToFile(rosField, filepath.Join(baseDir, "origin_map.bag"))
	if err != nil {
		fmt.Println("Error downloading origin map's bag file:", err)
		return "", "", err
	}
	c_log.GlobalLogger.Info("下载origin_map.bag文件到根目录 - 成功")

	// 创建压缩文件
	zipPath := filepath.Join(tmpDir, "simulationFile-"+id+".zip")
	zipFile, err := os.Create(zipPath)
	if err != nil {
		fmt.Println("Error creating ZIP file:", err)
		return "", "", err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()

	// 压缩文件夹
	if err := util.AddDirToZip(baseDir, zipWriter); err != nil {
		fmt.Println("Error adding directory to ZIP:", err)
		return "", "", err
	}
	fmt.Println("ZIP file created successfully.")

	c_log.GlobalLogger.Info("创建压缩文件 - 成功")

	return zipPath, tmpDir, nil
}

// CheckDataFileStatus 检查data目录是否存在
// @router /simulation/check/file/data/status [GET]
func CheckDataFileStatus(ctx context.Context, c *app.RequestContext) {
	sceneID := c.Query("id")
	fmt.Println("id", sceneID)

	// 根据id获取对应的oss文件列表
	allFileList, err := util.GetExactedMapFileById(sceneID)
	if err != nil {
		return
	}
	// 过滤特定后缀的文件列表（data.zip）
	simulationFileList := util.FilterBySuffixes(allFileList, config.DataFiltersuffixes...)
	if len(simulationFileList) == 0 {
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "data目录不存在"})
		return
	}
	c.JSON(consts.StatusOK, entity.HttpResult{Status: true, Code: "", Message: "data目录存在"})
}

// UploadPdfFile 将result.pdf文件上传到oss
// @router /simulation/upload/pdf [GET]
func UploadPdfFile(ctx context.Context, c *app.RequestContext) {
	equipmentNo := c.Query("equipmentNo")
	fmt.Println("equipmentNo", equipmentNo)

	sceneNo := c.Query("sceneNo")
	fmt.Println("sceneNo", sceneNo)

	timeStamp := c.Query("timeStamp")
	fmt.Println("timeStamp", timeStamp)

	round := c.Query("round")
	fmt.Println("round", round)

	header, err := c.FormFile("file")
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("get form err: %s", err.Error()))
		return
	}

	fileName := header.Filename
	fmt.Println("filename", fileName)

	ossObjectKey := config.SimulationOssBasePrefix + "/" + equipmentNo + "/" + sceneNo + "/" + timeStamp + "/" + round + "/" + fileName
	fmt.Println("ossObjectKey", ossObjectKey)

	f, _ := header.Open()
	defer f.Close()

	config.OssMutex.Lock()
	err = config.OssBucket.PutObject(ossObjectKey, f)
	config.OssMutex.Unlock()
	if err != nil {
		c_log.GlobalLogger.Error("程序异常退出。上传文件", fileName, "->", ossObjectKey, "出错：", err)
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "上传文件失败", Details: ""})
		return
	}
	c_log.GlobalLogger.Info("上传文件", fileName, "->", ossObjectKey, "成功。")
	c.JSON(consts.StatusOK, entity.HttpResult{Status: true, Code: "", Message: "上传文件成功", Details: ossObjectKey})
}

// UploadBagFile 将test-${i}.bag文件上传到oss
// @router /simulation/upload/bag [GET]
func UploadBagFile(ctx context.Context, c *app.RequestContext) {
	equipmentNo := c.Query("equipmentNo")
	fmt.Println("equipmentNo", equipmentNo)

	sceneNo := c.Query("sceneNo")
	fmt.Println("sceneNo", sceneNo)

	timeStamp := c.Query("timeStamp")
	fmt.Println("timeStamp", timeStamp)

	round := c.Query("round")
	fmt.Println("round", round)

	header, err := c.FormFile("file")
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("get form err: %s", err.Error()))
		return
	}

	fileName := header.Filename
	fmt.Println("filename", fileName)

	ossObjectKey := config.SimulationOssBasePrefix + "/" + equipmentNo + "/" + sceneNo + "/" + timeStamp + "/" + round + "/" + "test.bag"
	fmt.Println("ossObjectKey", ossObjectKey)

	f, _ := header.Open()
	defer f.Close()

	config.OssMutex.Lock()
	err = config.OssBucket.PutObject(ossObjectKey, f)
	config.OssMutex.Unlock()
	if err != nil {
		c_log.GlobalLogger.Error("程序异常退出。上传文件", fileName, "->", ossObjectKey, "出错：", err)
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "上传文件失败", Details: ""})
		return
	}
	c_log.GlobalLogger.Info("上传文件", fileName, "->", ossObjectKey, "成功。")
	c.JSON(consts.StatusOK, entity.HttpResult{Status: true, Code: "", Message: "上传文件成功", Details: ossObjectKey})
}

// UploadPgmFile 将map.pgm文件上传到oss
// @router /simulation/upload/pgm [GET]
func UploadPgmFile(ctx context.Context, c *app.RequestContext) {
	equipmentNo := c.Query("equipmentNo")
	fmt.Println("equipmentNo", equipmentNo)

	sceneNo := c.Query("sceneNo")
	fmt.Println("sceneNo", sceneNo)

	header, err := c.FormFile("file")
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("get form err: %s", err.Error()))
		return
	}

	fileName := header.Filename
	fmt.Println("filename", fileName)

	ossObjectKey := config.SimulationOssBasePrefix + "/" + equipmentNo + "/" + sceneNo + "/" + fileName
	fmt.Println("ossObjectKey", ossObjectKey)

	f, _ := header.Open()
	defer f.Close()

	config.OssMutex.Lock()
	err = config.OssBucket.PutObject(ossObjectKey, f)
	config.OssMutex.Unlock()
	if err != nil {
		c_log.GlobalLogger.Error("程序异常退出。上传文件", fileName, "->", ossObjectKey, "出错：", err)
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "上传文件失败", Details: ""})
		return
	}
	c_log.GlobalLogger.Info("上传文件", fileName, "->", ossObjectKey, "成功。")
	c.JSON(consts.StatusOK, entity.HttpResult{Status: true, Code: "", Message: "上传文件成功", Details: ossObjectKey})
}

// AddSimulationRecord 添加仿真测试记录
// @router /simulation/add/record [GET]
func AddSimulationRecord(ctx context.Context, c *app.RequestContext) {
	var records []*model.SimulationTestRecord
	err := c.BindAndValidate(&records)

	for _, record := range records {
		record.ID = uuid.NewV1().String()
	}
	fmt.Println("records", records)
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("get form err: %s", err.Error()))
		return
	}
	err = mysql.AddSimulationTestRecords(ctx, records)
	if err != nil {
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "仿真测试记录添加失败"})
		return
	}
	c.JSON(consts.StatusOK, entity.HttpResult{Status: true, Code: "", Message: "仿真测试记录添加成功"})
}

// QueryEvalRecord 查询算法评价记录（根据场景id及时间戳）
// @router /simulation/query/eval/record [GET]
func QueryEvalRecord(ctx context.Context, c *app.RequestContext) {
	sceneId := c.Query("sceneId")
	fmt.Println("sceneId", sceneId)

	testTime := c.Query("testTime")
	fmt.Println("testTime", testTime)

	records, err := mysql.QuerySimulationTestRecords(ctx, sceneId, testTime)
	if err != nil {
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "仿真测试记录查询失败", Details: ""})
		return
	}
	output, err := json.Marshal(records)
	if err != nil {
		c.JSON(consts.StatusOK, entity.HttpResult{Status: false, Code: "", Message: "仿真测试记录查询失败", Details: ""})
		return
	}
	c.JSON(consts.StatusOK, entity.HttpResult{Status: true, Code: "", Message: "仿真测试记录查询成功", Details: string(output)})
}

// DownloadFileByKeys 给定object key数组从oss下载（打包）文件
// @router /simulation/download/oss/key [POST]
func DownloadFileByKeys(ctx context.Context, c *app.RequestContext) {
	sceneId := c.Query("sceneId")
	fmt.Println("sceneId", sceneId)

	typeName := c.Query("typeName")
	fmt.Println("typeName", typeName)

	var req []string
	err := c.BindAndValidate(&req)
	if err != nil {
		c.String(consts.StatusBadRequest, err.Error())
		return
	}
	fmt.Println("req", req)
	if len(req) == 0 {
		c.JSON(consts.StatusBadRequest, entity.HttpResult{Status: false, Code: "", Message: "请求数据为空。"})
	} else if len(req) == 1 { // 只有一条数据则直接下载文件
		objectKey := req[0]
		// 从OSS下载文件
		reader, err := config.OssBucket.GetObject(objectKey)

		if err != nil {
			c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
			return
		}
		defer reader.Close()

		// 设置响应头
		c.Response.Header.Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filepath.Base(objectKey)))
		c.Response.Header.Set("Content-Type", "binary/octet-stream")

		// 将文件流式传输回客户端
		data, err := io.ReadAll(reader)
		if err != nil {
			panic(err)
		}
		if _, err := c.Write(data); err != nil {
			c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
			return
		}
	} else { // 多条数据先打包后下载文件
		fileName := sceneId + "_" + typeName + ".zip"
		filePath, tmpDir, err := generateZipByKey(ctx, req, fileName)
		if err != nil {
			c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
		}
		fmt.Println("filePath", filePath)

		// 检查文件是否存在
		if _, err := os.Stat(filePath); os.IsNotExist(err) {
			c.JSON(http.StatusNotFound, map[string]string{"error": "File not found"})
			return
		}

		// 打开文件
		f, err := os.Open(filePath)
		if err != nil {
			c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to open file"})
			return
		}
		defer f.Close()

		// 设置响应头
		c.Response.Header.Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filepath.Base(filePath)))
		c.Response.Header.Set("Content-Type", "binary/octet-stream")

		// 将文件流式传输回客户端
		data, err := io.ReadAll(f)
		if err != nil {
			panic(err)
		}
		if _, err := c.Write(data); err != nil {
			c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
			return
		}
		defer os.RemoveAll(tmpDir)
	}
}

// 根据object keys生成压缩包
func generateZipByKey(ctx context.Context, keys []string, fileName string) (file string, tmpDir string, err error) {
	// 创建临时文件夹
	tmpDir, err = os.MkdirTemp("", "temp-download-*")

	fmt.Println("tmpDir:", tmpDir)
	if err != nil {
		fmt.Println("Error creating temporary directory:", err)
		return "", "", err
	}
	c_log.GlobalLogger.Info("创建下载-临时文件夹：", tmpDir)

	// 创建根文件夹（文件打包的根目录）
	baseDir := filepath.Join(tmpDir, "data")
	if err := os.Mkdir(baseDir, 0755); err != nil {
		fmt.Println("Error creating subdirectory:", err)
		return "", "", err
	}
	c_log.GlobalLogger.Info("创建文件打包根目录：", baseDir)

	// 从oss下载文件到根目录
	for i, file := range keys {
		baseName := filepath.Base(file)
		ext := filepath.Ext(baseName)
		name := strings.TrimSuffix(baseName, ext)
		newName := name + "-" + strconv.Itoa(i+1) + ext
		err = config.OssBucket.GetObjectToFile(file, filepath.Join(baseDir, newName))
		if err != nil {
			fmt.Println("Error downloading file:", file, err)
			return "", "", err
		}
	}
	c_log.GlobalLogger.Info("下载oss文件到根目录 - 成功")

	// 创建压缩文件
	zipPath := filepath.Join(tmpDir, fileName)
	zipFile, err := os.Create(zipPath)
	if err != nil {
		fmt.Println("Error creating ZIP file:", err)
		return "", "", err
	}
	defer zipFile.Close()

	zipWriter := zip.NewWriter(zipFile)
	defer zipWriter.Close()

	// 压缩文件夹
	if err := util.AddDirToZip(baseDir, zipWriter); err != nil {
		fmt.Println("Error adding directory to ZIP:", err)
		return "", "", err
	}
	fmt.Println("ZIP file created successfully.")

	c_log.GlobalLogger.Info("创建压缩文件 - 成功")

	return zipPath, tmpDir, nil
}