kenaito-media-picker/internal/organize/organize.go

// Package organize 编排整理流程：探测 → 归类 → 收集 bucket → 一次性移动并记录去向。
package organize

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"kenaito-media-picker/internal/classify"
	"kenaito-media-picker/internal/digest"
	"kenaito-media-picker/internal/mover"
	"kenaito-media-picker/internal/probe"
	"kenaito-media-picker/internal/scanner"
	"kenaito-media-picker/internal/store"
)

// 特殊归类目录名
const (
	dirImagePSD     = "图片_PSD"
	dirImageErr     = "图片_读取异常"
	dirImageIgnored = "图片_已忽略"
	dirVideoErr     = "视频_读取异常"
	dirVideoIgnored = "视频_已忽略"
	dirAudio        = "音频"
)

// Organizer 整理器：按目标目录收集文件，处理完成后统一建目录并移动
type Organizer struct {
	pickerDir   string
	ffprobePath string
	store       *store.Store
	mover       *mover.Mover
	buckets     map[string][]string // 目标目录 → 文件路径
	order       []string            // 目标目录的创建顺序（保证输出稳定）
}

// New 创建整理器；store 可为 nil（数据库初始化失败时仅不记录信息）
func New(pickerDir, ffprobePath string, st *store.Store) *Organizer {
	return &Organizer{
		pickerDir:   pickerDir,
		ffprobePath: ffprobePath,
		store:       st,
		mover:       &mover.Mover{},
		buckets:     make(map[string][]string),
	}
}

// add 将文件归入某个子目录
func (o *Organizer) add(f scanner.File, subDir string) {
	dir := filepath.Join(o.pickerDir, subDir)
	if _, ok := o.buckets[dir]; !ok {
		o.order = append(o.order, dir)
	}
	o.buckets[dir] = append(o.buckets[dir], f.Path)
}

// record 将媒体信息写入数据库（大文件优先用指纹复用已入库的 MD5）
func (o *Organizer) record(f scanner.File, mediaType string, res *probe.ProbeResult) {
	st, err := os.Stat(f.Path)
	if err != nil {
		fmt.Printf("=== 读取文件信息失败: %s, 错误: %v\n", f.Path, err)
		return
	}
	size := st.Size()

	md5Str := ""
	headMd5, tailMd5 := "", ""
	if size > digest.FingerprintThreshold {
		if head, tail, ferr := digest.CalcFileFingerprint(f.Path, size); ferr == nil {
			headMd5, tailMd5 = head, tail
			if existing, ok := o.store.FindMediaByFingerprint(size, head, tail); ok && existing.FileMd5 != "" {
				md5Str = existing.FileMd5 // 指纹命中，复用MD5，跳过全量读取
			}
		}
	}
	if md5Str == "" {
		if md5Str, err = digest.CalcFileMd5(f.Path); err != nil {
			fmt.Printf("=== 计算MD5失败: %s, 错误: %v\n", f.Path, err)
			return
		}
	}

	codec := res.VideoCodec
	if mediaType == "audio" {
		codec = res.AudioCodec
	}
	info := &store.MediaInfo{
		FileMd5:   md5Str,
		FilePath:  f.Path,
		FileName:  f.Name,
		FileExt:   f.Ext,
		FileSize:  size,
		MimeType:  f.Mime,
		MediaType: mediaType,
		Width:     res.Width,
		Height:    res.Height,
		Duration:  res.Duration,
		Codec:     codec,
		Format:    res.Format,
		RawInfo:   res.RawJSON,
		HeadMd5:   headMd5,
		TailMd5:   tailMd5,
	}
	if err := o.store.SaveMedia(info); err != nil {
		fmt.Printf("=== 保存媒体信息失败: %s, 错误: %v\n", f.Path, err)
	}
}

// ProcessImages 处理全部图片：探测尺寸并归类
func (o *Organizer) ProcessImages(images []scanner.File) {
	success, fail, ignore := 0, 0, 0
	for _, f := range images {
		switch {
		case strings.EqualFold(f.Ext, ".psd"):
			o.add(f, dirImagePSD)
			success++
			continue
		case !probe.IsSupportImage(f.Ext):
			o.add(f, dirImageIgnored)
			ignore++
			continue
		}
		res, err := probe.ProbeImage(f.Path, o.ffprobePath)
		if err != nil {
			o.add(f, dirImageErr)
			fail++
			fmt.Printf("=== 异常图片: %s \n", f.Path)
			continue
		}
		o.record(f, "image", res)
		darkScore, lightScore := 0.0, 0.0
		if !strings.EqualFold(f.Ext, ".gif") {
			dark, light, berr := probe.AnalyzeBrightness(f.Path)
			if berr != nil {
				fmt.Printf("=== 明暗分析失败(按中性mid归类): %s \n", f.Path)
			} else {
				darkScore, lightScore = dark, light
			}
		}
		o.add(f, classify.ImageDir(f.Ext, res.Width, res.Height, darkScore, lightScore))
		success++
	}
	fmt.Printf("=== 图片处理完毕，成功: %d, 读取异常: %d, 忽略: %d \n\n", success, fail, ignore)
}

// ProcessMedia 处理媒体候选：ffprobe 判定视频/音频/不支持
func (o *Organizer) ProcessMedia(media []scanner.File) {
	video, audio, fail, ignore := 0, 0, 0, 0
	for _, f := range media {
		res, err := probe.ProbeMediaFile(o.ffprobePath, f.Path)
		if err != nil {
			// 仅 MIME 明确为视频/音频的文件才移入异常/忽略目录；
			// octet-stream/空 MIME 且 ffprobe 无法确认的原地忽略，不移动
			if strings.HasPrefix(f.Mime, "video/") {
				fail++
				o.add(f, dirVideoErr)
				fmt.Printf("=== 异常视频: %s \n", f.Path)
			} else if scanner.IsExplicitMediaMime(f.Mime) {
				ignore++
				o.add(f, dirVideoIgnored)
			} else {
				ignore++
				fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
			}
			continue
		}
		switch {
		case res.HasVideo:
			video++
			o.record(f, "video", res)
			o.add(f, classify.VideoDir(f.Ext, res.Width, res.Height))
		case res.HasAudio:
			audio++
			o.record(f, "audio", res)
			o.add(f, dirAudio)
		default:
			// ffprobe 无音视频流：MIME 明确为媒体的归入忽略目录，其余原地不动
			if scanner.IsExplicitMediaMime(f.Mime) {
				ignore++
				o.add(f, dirVideoIgnored)
			} else {
				ignore++
				fmt.Printf("=== 无法确认的非媒体文件, 原地忽略: %s \n", f.Path)
			}
		}
	}
	fmt.Printf("=== 媒体处理完毕，视频: %d, 音频: %d, 读取异常: %d, 忽略: %d \n\n", video, audio, fail, ignore)
}

// Flush 将所有收集到的文件移动到对应目录，并回填归类去向
func (o *Organizer) Flush() {
	for _, dir := range o.order {
		moved := o.mover.MoveToDir(o.buckets[dir], dir)
		for src, dest := range moved {
			if err := o.store.UpdateMediaDest(src, dest); err != nil {
				fmt.Printf("=== 更新媒体去向失败: %s, 错误: %v\n", src, err)
			}
		}
	}
}