shudao-main/shudao-go-backend/controllers/test.go

// Package controllers - test.go
//
// 📝 NOTE (说明)
// ================================================================================
// 本文件包含批量数据处理和测试工具函数。
// 这些函数仅用于开发和测试环境,不应在生产环境中调用。
//
// TODO: 建议将这些工具函数迁移到独立的CLI工具或脚本中。
// ================================================================================
package controllers

import (
	"fmt"
	"image"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"shudao-chat-go/models"
	"shudao-chat-go/utils"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// OSS配置信息 - 延迟初始化
var (
	testOssBucket    string
	testOssAccessKey string
	testOssSecretKey string
	testOssEndpoint  string
	testOssRegion    = "us-east-1"
	testOssInited    = false
)

func initTestOSSConfig() {
	if testOssInited {
		return
	}
	config := utils.GetOSSConfig()
	testOssBucket = config["bucket"]
	testOssAccessKey = config["access_key"]
	testOssSecretKey = config["secret_key"]
	testOssEndpoint = config["endpoint"]
	testOssInited = true
}

// 批量上传文件到OSS
func BatchUploadFilesToOSS() {
	initTestOSSConfig()
	fmt.Println("=== 开始批量上传文件到OSS ===")

	// 设置文件文件夹路径
	fileFolder := "C:/Users/allen/Desktop/分类文件/办法"
	fmt.Printf("目标文件夹: %s\n", fileFolder)

	// 检查文件夹是否存在
	if _, err := os.Stat(fileFolder); os.IsNotExist(err) {
		fmt.Printf("错误: 文件夹不存在: %s\n", fileFolder)
		return
	}

	fmt.Println("文件夹存在，开始扫描...")

	// 获取所有文件
	files, err := getDocumentFiles(fileFolder)
	if err != nil {
		fmt.Printf("错误: 获取文件失败: %v\n", err)
		return
	}

	fmt.Printf("找到 %d 个文件\n", len(files))

	// 创建S3会话
	sess, err := createS3Session()
	if err != nil {
		fmt.Printf("错误: 创建S3会话失败: %v\n", err)
		return
	}
	s3Client := s3.New(sess)

	// 处理每个文件
	successCount := 0
	errorCount := 0

	for i, filePath := range files {
		fmt.Printf("处理文件 %d/%d: %s\n", i+1, len(files), filepath.Base(filePath))

		// 获取文件名（不含路径）
		fileName := filepath.Base(filePath)

		// 上传文件到OSS
		fileURL, err := uploadFileToOSSTest(s3Client, filePath, fileName)
		if err != nil {
			fmt.Printf("  错误: 上传文件失败 - %v\n", err)
			errorCount++
			continue
		}

		fmt.Printf("  文件上传成功: %s\n", fileURL)
		//去除Filename点后面的内容
		fileName = strings.TrimSuffix(fileName, filepath.Ext(fileName))
		//去掉fileName前面“（内部）”这两个字
		fileName = strings.TrimPrefix(fileName, "（内部）")

		// 保存到index_file表
		indexFile := models.PolicyFile{
			PolicyName:       fileName,
			PolicyFileUrl:    fileURL,
			FileType:         0,
			FileTag:          "内部法规",
			PublishTime:      time.Now().Unix(),
			PolicyType:       4,
			PolicyDepartment: "内部法规",
			ViewCount:        0,
			PolicyContent:    "内部法规",
		}

		result := models.DB.Create(&indexFile)
		if result.Error != nil {
			fmt.Printf("  错误: 保存到数据库失败 - %v\n", result.Error)
			errorCount++
			continue
		}

		fmt.Printf("  数据库保存成功，ID: %d\n", indexFile.ID)
		successCount++

		// 添加延迟避免请求过于频繁
		// time.Sleep(100 * time.Millisecond)
	}

	fmt.Println("处理完成！")
	fmt.Printf("成功处理: %d 个文件\n", successCount)
	fmt.Printf("处理失败: %d 个文件\n", errorCount)
}

// 批量匹配图片
func BatchMatchImages() {
	fmt.Println("=== 开始批量匹配图片 ===")

	// 获取test表中所有有title1的记录
	var testRecords []models.Test
	result := models.DB.Where("title1 != '' AND title1 IS NOT NULL").Find(&testRecords)
	if result.Error != nil {
		fmt.Printf("错误: 获取test表记录失败: %v\n", result.Error)
		return
	}

	fmt.Printf("找到 %d 条test表记录\n", len(testRecords))

	successCount := 0
	errorCount := 0

	for i, testRecord := range testRecords {
		fmt.Printf("处理记录 %d/%d: ID=%d, Title1=%s\n", i+1, len(testRecords), testRecord.ID, testRecord.Title1)

		// 根据title1在third_scene表中查找匹配的记录
		var thirdScene models.ThirdScene
		result := models.DB.Where("third_scene_name = ?", testRecord.Title1).First(&thirdScene)
		if result.Error != nil {
			fmt.Printf("  错误: 在third_scene表中未找到匹配记录 - %v\n", result.Error)
			errorCount++
			continue
		}

		fmt.Printf("  找到匹配的third_scene记录: ID=%d, Name=%s\n", thirdScene.ID, thirdScene.ThirdSceneName)

		// 更新third_scene表的correct_example_image和wrong_example_image
		updateData := map[string]interface{}{
			"correct_example_image": testRecord.Title2,
			"wrong_example_image":   testRecord.Title3,
		}

		result = models.DB.Model(&thirdScene).Updates(updateData)
		if result.Error != nil {
			fmt.Printf("  错误: 更新third_scene表失败 - %v\n", result.Error)
			errorCount++
			continue
		}

		fmt.Printf("  更新成功: correct_example_image=%s, wrong_example_image=%s\n", testRecord.Title2, testRecord.Title3)
		successCount++
	}

	fmt.Println("匹配完成！")
	fmt.Printf("成功匹配: %d 条记录\n", successCount)
	fmt.Printf("匹配失败: %d 条记录\n", errorCount)
}

// BatchUploadImages 批量上传图片函数
func BatchUploadImages() {
	fmt.Println("=== 开始批量上传图片 ===")

	// 设置图片文件夹路径
	imageFolder := "C:/Users/allen/Desktop/隐患识别图/错误图片"
	fmt.Printf("目标文件夹: %s\n", imageFolder)

	// 检查文件夹是否存在
	if _, err := os.Stat(imageFolder); os.IsNotExist(err) {
		fmt.Printf("错误: 文件夹不存在: %s\n", imageFolder)
		return
	}

	fmt.Println("文件夹存在，开始扫描...")

	// 获取所有图片文件
	imageFiles, err := getImageFiles(imageFolder)
	if err != nil {
		fmt.Printf("错误: 获取图片文件失败: %v\n", err)
		return
	}

	fmt.Printf("找到 %d 个图片文件\n", len(imageFiles))

	// 创建S3会话
	sess, err := createS3Session()
	if err != nil {
		fmt.Printf("错误: 创建S3会话失败: %v\n", err)
		return
	}
	s3Client := s3.New(sess)

	// 处理每个图片文件（只处理前10个文件进行测试）
	successCount := 0
	errorCount := 0

	// 限制处理文件数量

	for i, imageFile := range imageFiles {
		fmt.Printf("处理文件 %d/%d: %s\n", i+1, len(imageFiles), filepath.Base(imageFile))

		// 提取文件名中的序号
		fileID, err := extractIDFromFilename(filepath.Base(imageFile))
		if err != nil {
			fmt.Printf("  错误: 无法提取序号 - %v\n", err)
			errorCount++
			continue
		}

		fmt.Printf("  提取的序号: %d\n", fileID)

		// 检查数据库中是否存在对应的记录
		var testRecord models.Test
		result := models.DB.First(&testRecord, fileID)
		if result.Error != nil {
			fmt.Printf("  错误: 数据库中未找到ID为 %d 的记录 - %v\n", fileID, result.Error)
			errorCount++
			continue
		}

		fmt.Printf("  找到数据库记录: ID=%d, Title1=%s\n", testRecord.ID, testRecord.Title1)

		// 上传图片到OSS
		imageURL, err := uploadImageToOSSTest(s3Client, imageFile, fileID)
		if err != nil {
			fmt.Printf("  错误: 上传图片失败 - %v\n", err)
			errorCount++
			continue
		}

		fmt.Printf("  图片上传成功: %s\n", imageURL)

		// 更新数据库记录
		result = models.DB.Model(&testRecord).Update("title3", imageURL)
		if result.Error != nil {
			fmt.Printf("  错误: 更新数据库失败 - %v\n", result.Error)
			errorCount++
			continue
		}

		fmt.Printf("  数据库更新成功\n")
		successCount++

		// 添加延迟避免请求过于频繁
		time.Sleep(100 * time.Millisecond)
	}

	fmt.Println("处理完成！")
	fmt.Printf("成功处理: %d 个文件\n", successCount)
	fmt.Printf("处理失败: %d 个文件\n", errorCount)
}

// 获取文件夹中的所有文档文件
func getDocumentFiles(folderPath string) ([]string, error) {
	var documentFiles []string

	// 支持的文档格式
	documentExts := map[string]bool{
		".pdf":  true,
		".doc":  true,
		".docx": true,
		".txt":  true,
		".rtf":  true,
		".odt":  true,
		".xls":  true,
		".xlsx": true,
		".ppt":  true,
		".pptx": true,
	}

	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() {
			ext := strings.ToLower(filepath.Ext(path))
			if documentExts[ext] {
				documentFiles = append(documentFiles, path)
			}
		}

		return nil
	})

	return documentFiles, err
}

// 获取文件夹中的所有图片文件
func getImageFiles(folderPath string) ([]string, error) {
	var imageFiles []string

	// 支持的图片格式
	imageExts := map[string]bool{
		".jpg":  true,
		".jpeg": true,
		".png":  true,
		".gif":  true,
		".bmp":  true,
		".webp": true,
		".tiff": true,
		".svg":  true,
		".ico":  true,
	}

	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() {
			ext := strings.ToLower(filepath.Ext(path))
			if imageExts[ext] {
				imageFiles = append(imageFiles, path)
			}
		}

		return nil
	})

	return imageFiles, err
}

// 从文件名中提取序号
func extractIDFromFilename(filename string) (int, error) {
	// 移除文件扩展名
	nameWithoutExt := strings.TrimSuffix(filename, filepath.Ext(filename))

	// 使用正则表达式提取开头的数字
	re := regexp.MustCompile(`^(\d+)`)
	matches := re.FindStringSubmatch(nameWithoutExt)

	if len(matches) < 2 {
		return 0, fmt.Errorf("文件名中未找到数字序号: %s", filename)
	}

	id, err := strconv.Atoi(matches[1])
	if err != nil {
		return 0, fmt.Errorf("无法解析序号: %s", matches[1])
	}

	return id, nil
}

// 创建S3会话
func createS3Session() (*session.Session, error) {
	initTestOSSConfig()
	s3Config := &aws.Config{
		Credentials:      credentials.NewStaticCredentials(testOssAccessKey, testOssSecretKey, ""),
		Endpoint:         aws.String(testOssEndpoint),
		Region:           aws.String(testOssRegion),
		S3ForcePathStyle: aws.Bool(true),
	}

	sess, err := session.NewSession(s3Config)
	if err != nil {
		return nil, err
	}

	// 验证凭据
	_, err = sess.Config.Credentials.Get()
	if err != nil {
		return nil, fmt.Errorf("凭据验证失败: %v", err)
	}

	return sess, nil
}

// 上传文件到OSS
func uploadFileToOSSTest(s3Client *s3.S3, filePath string, fileName string) (string, error) {
	// 打开文件
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 生成文件名
	ext := filepath.Ext(filePath)
	ossFileName := fmt.Sprintf("documents/%d_%s",
		time.Now().Unix(), fileName)

	// 读取文件内容
	fileBytes, err := io.ReadAll(file)
	if err != nil {
		return "", err
	}

	// 确定Content-Type
	contentType := getContentType(ext)

	// 上传到S3
	_, err = s3Client.PutObject(&s3.PutObjectInput{
		Bucket:      aws.String(testOssBucket),
		Key:         aws.String(ossFileName),
		Body:        aws.ReadSeekCloser(strings.NewReader(string(fileBytes))),
		ContentType: aws.String(contentType),
		ACL:         aws.String("public-read"),
	})

	if err != nil {
		return "", err
	}

	// 生成访问URL
	fileURL := fmt.Sprintf("%s/%s/%s", testOssEndpoint, testOssBucket, ossFileName)
	return fileURL, nil
}

// 根据文件扩展名获取Content-Type
func getContentType(ext string) string {
	contentTypes := map[string]string{
		".pdf":  "application/pdf",
		".doc":  "application/msword",
		".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
		".txt":  "text/plain",
		".rtf":  "application/rtf",
		".odt":  "application/vnd.oasis.opendocument.text",
		".xls":  "application/vnd.ms-excel",
		".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
		".ppt":  "application/vnd.ms-powerpoint",
		".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
	}

	if contentType, exists := contentTypes[strings.ToLower(ext)]; exists {
		return contentType
	}
	return "application/octet-stream" // 默认类型
}

// 上传图片到OSS
func uploadImageToOSSTest(s3Client *s3.S3, imagePath string, fileID int) (string, error) {
	// 打开图片文件
	file, err := os.Open(imagePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 生成文件名
	ext := filepath.Ext(imagePath)
	fileName := fmt.Sprintf("batch_upload/%d_%d%s",
		time.Now().Unix(), fileID, ext)

	// 读取文件内容
	fileBytes, err := io.ReadAll(file)
	if err != nil {
		return "", err
	}

	// 上传到S3
	_, err = s3Client.PutObject(&s3.PutObjectInput{
		Bucket: aws.String(testOssBucket),
		Key:    aws.String(fileName),
		Body:   aws.ReadSeekCloser(strings.NewReader(string(fileBytes))),
		ACL:    aws.String("public-read"),
	})

	if err != nil {
		return "", err
	}

	// 生成访问URL
	imageURL := fmt.Sprintf("%s/%s/%s", testOssEndpoint, testOssBucket, fileName)
	return imageURL, nil
}

// BatchCompressAndReuploadImages 批量压缩和重新上传third_scene表中的图片
func BatchCompressAndReuploadImages() {
	fmt.Println("=== 开始批量压缩和重新上传图片 ===")

	// 获取third_scene表中所有记录
	var thirdScenes []models.ThirdScene
	result := models.DB.Find(&thirdScenes)
	if result.Error != nil {
		fmt.Printf("错误: 获取third_scene表记录失败: %v\n", result.Error)
		return
	}

	fmt.Printf("找到 %d 条third_scene记录\n", len(thirdScenes))

	// 创建S3会话
	sess, err := createS3Session()
	if err != nil {
		fmt.Printf("错误: 创建S3会话失败: %v\n", err)
		return
	}
	s3Client := s3.New(sess)

	successCount := 0
	errorCount := 0
	skipCount := 0

	for i, thirdScene := range thirdScenes {
		fmt.Printf("处理记录 %d/%d: ID=%d, Name=%s\n", i+1, len(thirdScenes), thirdScene.ID, thirdScene.ThirdSceneName)

		// 处理correct_example_image
		if thirdScene.CorrectExampleImage != "" {
			fmt.Printf("  处理correct_example_image: %s\n", thirdScene.CorrectExampleImage)
			newURL, err := compressAndReuploadImage(s3Client, thirdScene.CorrectExampleImage, fmt.Sprintf("correct_%d", thirdScene.ID))
			if err != nil {
				fmt.Printf("  错误: 处理correct_example_image失败 - %v\n", err)
				errorCount++
			} else {
				thirdScene.CorrectExampleImage = newURL
				fmt.Printf("  correct_example_image更新成功: %s\n", newURL)
			}
		} else {
			fmt.Printf("  correct_example_image为空，跳过\n")
			skipCount++
		}

		// 处理wrong_example_image
		if thirdScene.WrongExampleImage != "" {
			fmt.Printf("  处理wrong_example_image: %s\n", thirdScene.WrongExampleImage)
			newURL, err := compressAndReuploadImage(s3Client, thirdScene.WrongExampleImage, fmt.Sprintf("wrong_%d", thirdScene.ID))
			if err != nil {
				fmt.Printf("  错误: 处理wrong_example_image失败 - %v\n", err)
				errorCount++
			} else {
				thirdScene.WrongExampleImage = newURL
				fmt.Printf("  wrong_example_image更新成功: %s\n", newURL)
			}
		} else {
			fmt.Printf("  wrong_example_image为空，跳过\n")
			skipCount++
		}

		// 更新数据库记录
		result := models.DB.Model(&thirdScene).Updates(map[string]interface{}{
			"correct_example_image": thirdScene.CorrectExampleImage,
			"wrong_example_image":   thirdScene.WrongExampleImage,
		})
		if result.Error != nil {
			fmt.Printf("  错误: 更新数据库失败 - %v\n", result.Error)
			errorCount++
		} else {
			fmt.Printf("  数据库更新成功\n")
			successCount++
		}

		// 添加延迟避免请求过于频繁
		time.Sleep(200 * time.Millisecond)
	}

	fmt.Println("处理完成！")
	fmt.Printf("成功处理: %d 条记录\n", successCount)
	fmt.Printf("处理失败: %d 条记录\n", errorCount)
	fmt.Printf("跳过空字段: %d 次\n", skipCount)
}

// compressAndReuploadImage 压缩并重新上传图片
func compressAndReuploadImage(s3Client *s3.S3, imageURL string, prefix string) (string, error) {
	// 下载图片
	tempFilePath, err := downloadImage(imageURL)
	if err != nil {
		return "", fmt.Errorf("下载图片失败: %v", err)
	}
	defer os.Remove(tempFilePath) // 清理临时文件

	// 压缩图片
	compressedFilePath, err := compressImageFile(tempFilePath)
	if err != nil {
		return "", fmt.Errorf("压缩图片失败: %v", err)
	}
	defer os.Remove(compressedFilePath) // 清理临时文件

	// 上传压缩后的图片
	newURL, err := uploadCompressedImageToOSS(s3Client, compressedFilePath, prefix)
	if err != nil {
		return "", fmt.Errorf("上传压缩图片失败: %v", err)
	}

	return newURL, nil
}

// downloadImage 下载图片到临时文件
func downloadImage(imageURL string) (string, error) {
	// 创建HTTP请求
	resp, err := http.Get(imageURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("下载失败，状态码: %d", resp.StatusCode)
	}

	// 创建临时文件
	tempFile, err := os.CreateTemp("", "download_*.jpg")
	if err != nil {
		return "", err
	}
	defer tempFile.Close()

	// 复制内容到临时文件
	_, err = io.Copy(tempFile, resp.Body)
	if err != nil {
		os.Remove(tempFile.Name())
		return "", err
	}

	return tempFile.Name(), nil
}

// compressImageFile 压缩图片到1M以下
func compressImageFile(inputPath string) (string, error) {
	// 打开原始图片
	file, err := os.Open(inputPath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 解码图片
	img, format, err := image.Decode(file)
	if err != nil {
		return "", err
	}

	// 创建临时文件用于压缩后的图片
	tempFile, err := os.CreateTemp("", "compressed_*.jpg")
	if err != nil {
		return "", err
	}
	defer tempFile.Close()

	// 尝试不同的质量级别直到文件大小小于1MB
	targetSize := int64(1024 * 1024) // 1MB
	quality := 90

	for quality > 10 {
		// 重置文件指针
		tempFile.Seek(0, 0)
		tempFile.Truncate(0)

		// 根据格式编码图片
		if format == "png" {
			err = png.Encode(tempFile, img)
		} else {
			err = jpeg.Encode(tempFile, img, &jpeg.Options{Quality: quality})
		}

		if err != nil {
			return "", err
		}

		// 检查文件大小
		fileInfo, err := tempFile.Stat()
		if err != nil {
			return "", err
		}

		if fileInfo.Size() <= targetSize {
			break
		}

		quality -= 10
	}

	// 如果还是太大，尝试缩放图片
	if quality <= 10 {
		// 获取图片尺寸
		bounds := img.Bounds()
		width := bounds.Dx()
		height := bounds.Dy()

		// 计算缩放比例
		scale := 0.8
		newWidth := int(float64(width) * scale)
		newHeight := int(float64(height) * scale)

		// 创建缩放后的图片
		resizedImg := resizeImageFile(img, newWidth, newHeight)

		// 重置文件指针
		tempFile.Seek(0, 0)
		tempFile.Truncate(0)

		// 编码缩放后的图片
		err = jpeg.Encode(tempFile, resizedImg, &jpeg.Options{Quality: 80})
		if err != nil {
			return "", err
		}
	}

	return tempFile.Name(), nil
}

// resizeImageFile 缩放图片
func resizeImageFile(img image.Image, width, height int) image.Image {
	// 简单的最近邻缩放
	bounds := img.Bounds()
	srcWidth := bounds.Dx()
	srcHeight := bounds.Dy()

	// 创建新的图片
	newImg := image.NewRGBA(image.Rect(0, 0, width, height))

	// 计算缩放比例
	xRatio := float64(srcWidth) / float64(width)
	yRatio := float64(srcHeight) / float64(height)

	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			srcX := int(float64(x) * xRatio)
			srcY := int(float64(y) * yRatio)
			newImg.Set(x, y, img.At(srcX, srcY))
		}
	}

	return newImg
}

// uploadCompressedImageToOSS 上传压缩后的图片到OSS
func uploadCompressedImageToOSS(s3Client *s3.S3, imagePath string, prefix string) (string, error) {
	// 打开图片文件
	file, err := os.Open(imagePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 生成文件名
	fileName := fmt.Sprintf("compressed_images/%s_%d.jpg", prefix, time.Now().Unix())

	// 读取文件内容
	fileBytes, err := io.ReadAll(file)
	if err != nil {
		return "", err
	}

	// 上传到S3
	_, err = s3Client.PutObject(&s3.PutObjectInput{
		Bucket:      aws.String(testOssBucket),
		Key:         aws.String(fileName),
		Body:        aws.ReadSeekCloser(strings.NewReader(string(fileBytes))),
		ContentType: aws.String("image/jpeg"),
		ACL:         aws.String("public-read"),
	})

	if err != nil {
		return "", err
	}

	// 生成访问URL
	imageURL := fmt.Sprintf("%s/%s/%s", testOssEndpoint, testOssBucket, fileName)
	return imageURL, nil
}