fonchain-fiee/pkg/utils/launch/launch.go

package launch

import (
	"archive/zip"
	"fmt"
	"github.com/fonchain_enterprise/fonchain-main/pkg/request"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// 下载文件并存储在指定文件夹
func downloadFile(url string, folder string, fileName string, wg *sync.WaitGroup, fileChan chan<- string, errChan chan<- error) {
	defer wg.Done()

	// 创建文件夹
	err := os.MkdirAll(folder, os.ModePerm)
	if err != nil {
		errChan <- err
		return
	}
	// 获取文件名
	//fileName := filepath.Base(url)
	filePath := filepath.Join(folder, fileName)

	// 创建文件
	out, err := os.Create(filePath)
	if err != nil {
		errChan <- err
		return
	}
	defer out.Close()

	// 下载文件
	resp, err := http.Get(url)
	if err != nil {
		errChan <- err
		return
	}
	defer resp.Body.Close()

	// 写入文件
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		errChan <- err
		return
	}

	fileChan <- filePath
}

// 将文件打包成ZIP
func zipFiles(zipFileName string, files map[string][]string) error {
	newZipFile, err := os.Create(zipFileName)
	if err != nil {
		return err
	}
	defer newZipFile.Close()

	zipWriter := zip.NewWriter(newZipFile)

	for folder, fileList := range files {
		for _, file := range fileList {
			if err := addFileToZip(zipWriter, file, folder); err != nil {
				return err
			}
		}
	}

	err = zipWriter.Close()
	if err != nil {
		return err
	}

	return nil
}

func addFileToZip(zipWriter *zip.Writer, filename string, folder string) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	// 获取文件信息
	info, err := file.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}

	// 使用文件夹路径
	header.Name = filepath.Join(folder, filepath.Base(filename))
	//header.Name = filepath.Base(filename)

	writer, err := zipWriter.CreateHeader(header)
	if err != nil {
		return err
	}

	_, err = io.Copy(writer, file)
	return err
}

// 删除目录及其内容
func removeDir(dir string) error {
	return os.RemoveAll(dir)
}

// 处理下载和压缩请求
func DownloadAndZip(reqData request.Response, fileType string, w http.ResponseWriter) {
	// 标记是否只包含一个 attachment
	isSingleAttachment := false
	if len(reqData.Contracts) == 1 && len(reqData.Contracts[0].Attachment) == 1 {
		isSingleAttachment = true
	}

	if isSingleAttachment {
		// 只包含一个 contract，且该 contract 只有一个 attachment，则直接下载该文件
		attachment := reqData.Contracts[0].Attachment[0]
		if attachment.URL == "" {
			http.Error(w, "URL not provided", http.StatusBadRequest)
			return
		}

		// 下载文件
		resp, err := http.Get(attachment.URL)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to download file: %v", err), http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		// 设置响应头
		w.Header().Set("Content-Type", resp.Header.Get("Content-Type"))
		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", attachment.Name))

		// 将文件内容发送给客户端
		_, err = io.Copy(w, resp.Body)
		if err != nil {
			log.Printf("Failed to send file content: %v", err)
		}

		return
	}

	// 多个 URL，则创建 ZIP 文件
	tempDir := fmt.Sprintf("temp_%d", time.Now().UnixNano())
	os.MkdirAll(tempDir, os.ModePerm)

	filesToZip := make(map[string][]string)
	var wg sync.WaitGroup
	fileChan := make(chan string)
	errChan := make(chan error)
	//fileNameCount := make(map[string]int)
	fileNameCount := make(map[string]map[string]int) // 存储每个附件下的文件名计数器
	for _, contract := range reqData.Contracts {
		contractFolder := contract.Theme
		for _, attachment := range contract.Attachment {
			wg.Add(1)
			fileName := attachment.Name
			ext := filepath.Ext(fileName)
			name := strings.TrimSuffix(fileName, ext)
			// 检查文件名是否重复，如果重复则对文件名进行枚举，否则不需要
			if counts, ok := fileNameCount[contractFolder]; ok {
				if count, exists := counts[fileName]; exists {
					counts[fileName]++
					fileName = fmt.Sprintf("%s(%d)%s", name, count, ext)
				} else {
					counts[fileName] = 1
				}
			} else {
				fileNameCount[contractFolder] = map[string]int{fileName: 1}
			}
			go downloadFile(attachment.URL, tempDir, fileName, &wg, fileChan, errChan)
			filesToZip[contractFolder] = append(filesToZip[contractFolder], filepath.Join(tempDir, fileName))
		}
	}

	// Goroutine to handle file paths from fileChan
	go func() {
		for filePath := range fileChan {
			log.Printf("Downloaded: %s", filePath)
		}
	}()

	// Goroutine to handle errors from errChan
	go func() {
		for err := range errChan {
			log.Printf("Error: %v", err)
		}
	}()

	wg.Wait()
	close(fileChan)
	close(errChan)

	// 创建ZIP文件
	zipFileName := filepath.Join(tempDir, fileType+".zip")
	err := zipFiles(zipFileName, filesToZip)
	if err != nil {
		log.Printf("Failed to create zip file: %v", err)
		http.Error(w, "Failed to create zip file", http.StatusInternalServerError)
		return
	}
	zipFile, err := os.Open(zipFileName)
	if err != nil {
		log.Printf("Failed to open zip file: %v", err)
		http.Error(w, "Failed to open zip file", http.StatusInternalServerError)
		return
	}
	// 返回ZIP文件
	w.Header().Set("Content-Type", "application/zip")
	encodedFileName := url.PathEscape(fmt.Sprintf("%s.zip", fileType))
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename*=UTF-8''%s", encodedFileName))
	io.Copy(w, zipFile)
	// 删除临时文件和目录
	defer func() {
		// 关闭文件句柄
		if err := zipFile.Close(); err != nil {
			log.Printf("Failed to close zip file: %v", err)
		}
		// 删除 ZIP 文件
		if err := os.Remove(zipFileName); err != nil {
			log.Printf("Failed to delete zip file: %v", err)
		}
		// 删除临时目录及其内容
		if err := removeDir(tempDir); err != nil {
			log.Printf("Failed to delete temp directory: %v", err)
		}
	}()
}