ai_scheduler/internal/biz/handle/file.go

package handle

import (
	"ai_scheduler/internal/data/constants"
	"ai_scheduler/internal/entitys"
	"ai_scheduler/internal/pkg/l_request"
	"bytes"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"

	"github.com/gabriel-vasile/mimetype"
)

// HandleRecognizeFile  这里的目的是无论将什么类型的file都转为二进制格式
// 最终输出：1.将 files.FileData 填充为文件的二进制数据 2.将 files.FileType 填充为文件的类型（当前为 constants.Caller，兼容写入其字符串值）
// 判断文件大小（统一限制为10MB）；判断文件类型；判断文件是否合法（类型在白名单映射中）；无法识别/非法/超限→填充unknown并兼容返回
// 若 FileData 不存在 且 FileUrl 不存在， 则直接退出
// 若 FileData 存在 FileType 存在， 则直接退出
// 若 FileData 存在 FileType 不存在， 则根据 FileData 推断文件类型并填充 FileType
// 若 FileUrl 存在， 则下载文件并填充 FileData 和 FileType
func HandleRecognizeFile(files *entitys.RecognizeFile) {
	if files == nil {
		return
	}

	const maxSize = 10 * 1024 * 1024 // 10MB 上限

	// 工具：根据 MIME 或扩展名映射到 FileType
	mapToFileType := func(s string) constants.FileType {
		if len(s) == 0 {
			return constants.FileTypeUnknown
		}
		s = strings.ToLower(strings.TrimSpace(s))
		for ft, items := range constants.FileTypeMappings {
			for _, item := range items {
				if !strings.HasPrefix(item, ".") { // MIME
					if s == item {
						return ft
					}
				} else { // 扩展名
					if s == item {
						return ft
					}
				}
			}
		}
		return constants.FileTypeUnknown
	}

	// 分支1：无数据、无URL→直接返回
	if len(files.FileData) == 0 && len(files.FileUrl) == 0 {
		return
	}

	// 分支2：已有数据且已有类型→直接返回
	if len(files.FileData) > 0 && len(strings.TrimSpace(files.FileType.String())) > 0 {
		return
	}

	// 分支3：仅有数据、无类型→内容检测并填充
	if len(files.FileData) > 0 && len(strings.TrimSpace(files.FileType.String())) == 0 {
		if len(files.FileData) > maxSize {
			files.FileType = constants.Caller(constants.FileTypeUnknown)
			return
		}

		reader := bytes.NewReader(files.FileData)
		detected := detectFileType(reader, "")
		if detected == constants.FileTypeUnknown {
			files.FileType = constants.Caller(constants.FileTypeUnknown)
			return
		}
		files.FileType = constants.Caller(detected)
		return
	}

	// 分支4：存在URL→下载并填充数据与类型
	if len(files.FileUrl) > 0 {
		fileBytes, contentType, err := downloadFile(files.FileUrl)
		if err != nil || len(fileBytes) == 0 {
			files.FileType = constants.Caller(constants.FileTypeUnknown)
			return
		}

		if len(fileBytes) > maxSize {
			// 超限：不写入数据，类型置 unknown
			files.FileType = constants.Caller(constants.FileTypeUnknown)
			return
		}

		// 优先使用响应头的 Content-Type 映射
		detected := mapToFileType(contentType)

		if detected == constants.FileTypeUnknown {
			// 回退：内容检测 + URL 文件名扩展名辅助
			var fname string
			if u, perr := url.Parse(files.FileUrl); perr == nil {
				fname = filepath.Base(u.Path)
			}
			reader := bytes.NewReader(fileBytes)
			detected = detectFileType(reader, fname)
		}

		// 写入数据
		files.FileData = fileBytes

		if detected == constants.FileTypeUnknown {
			files.FileType = constants.Caller(constants.FileTypeUnknown)
			return
		}
		files.FileType = constants.Caller(detected)
		return
	}
}

// 下载文件并返回二进制数据、MIME 类型
func downloadFile(fileUrl string) (fileBytes []byte, contentType string, err error) {
	if len(fileUrl) == 0 {
		return
	}
	req := l_request.Request{
		Method: "GET",
		Url:    fileUrl,
		Headers: map[string]string{
			"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
			"Accept":     "image/webp,image/apng,image/*,*/*;q=0.8",
		},
	}
	res, err := req.Send()
	if err != nil {
		return
	}
	var ex bool
	if contentType, ex = res.Headers["Content-Type"]; !ex {
		err = errors.New("Content-Type不存在")
		return
	}

	if res.StatusCode != http.StatusOK {
		err = fmt.Errorf("server returned non-200 status: %d", res.StatusCode)
	}
	fileBytes = res.Content

	return fileBytes, contentType, nil
}

// detectFileType 判断文件类型
func detectFileType(file io.ReadSeeker, filename string) constants.FileType {
	// 1. 读取文件头检测 MIME
	buffer := make([]byte, 512)
	n, _ := file.Read(buffer)
	file.Seek(0, io.SeekStart) // 重置读取位置

	detectedMIME := mimetype.Detect(buffer[:n]).String()
	for fileType, items := range constants.FileTypeMappings {
		for _, item := range items {
			if !strings.HasPrefix(item, ".") && item == detectedMIME {
				return fileType
			}
		}
	}

	// 2. 备用：通过扩展名检测
	ext := strings.ToLower(filepath.Ext(filename))
	for fileType, items := range constants.FileTypeMappings {
		for _, item := range items {
			if strings.HasPrefix(item, ".") && item == ext {
				return fileType
			}
		}
	}

	return constants.FileTypeUnknown
}