handler/task.go

package handler

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"video-inpaint/config"
	"video-inpaint/service"

	"github.com/gin-gonic/gin"
)

type TaskHandler struct{}

func (h *TaskHandler) Submit(c *gin.Context) {
	// 解析水印区域列表
	regionsStr := c.PostForm("regions")
	if regionsStr == "" {
		Fail(c, "regions 参数不能为空")
		return
	}
	var regions []service.Region
	if err := json.Unmarshal([]byte(regionsStr), &regions); err != nil {
		Fail(c, "regions 格式错误，需为 JSON 数组")
		return
	}
	if len(regions) == 0 {
		Fail(c, "至少需要一个水印区域")
		return
	}
	for i, r := range regions {
		if r.W <= 0 || r.H <= 0 {
			Fail(c, fmt.Sprintf("区域 %d 宽高必须大于 0", i))
			return
		}
	}

	// 获取上传文件
	file, err := c.FormFile("file")
	if err != nil {
		Fail(c, "未收到视频文件")
		return
	}

	maxMB := int64(500)
	if config.Cfg != nil && config.Cfg.Video.MaxUploadMB > 0 {
		maxMB = config.Cfg.Video.MaxUploadMB
	}
	if file.Size > maxMB*1024*1024 {
		Fail(c, fmt.Sprintf("文件过大，最大支持 %dMB", maxMB))
		return
	}

	// 创建任务并确定存储路径
	task := service.NewTask()

	tempDir := "./temp"
	if config.Cfg != nil && config.Cfg.Video.TempDir != "" {
		tempDir = config.Cfg.Video.TempDir
	}
	taskDir := filepath.Join(tempDir, task.TaskID)
	if err := os.MkdirAll(taskDir, 0755); err != nil {
		Fail(c, "创建任务目录失败")
		return
	}

	ext := filepath.Ext(file.Filename)
	if ext == "" {
		ext = ".mp4"
	}
	inputPath := filepath.Join(taskDir, "input"+ext)

	if err := c.SaveUploadedFile(file, inputPath); err != nil {
		os.RemoveAll(taskDir)
		Fail(c, "保存视频失败")
		return
	}

	// 异步启动处理
	service.ProcessVideo(task, inputPath, regions)

	Success(c, "ok", gin.H{"task_id": task.TaskID})
}

func (h *TaskHandler) Progress(c *gin.Context) {
	taskID := c.Query("task_id")
	if taskID == "" {
		Fail(c, "task_id 不能为空")
		return
	}

	task, ok := service.GetTask(taskID)
	if !ok {
		Fail(c, "任务不存在或已过期")
		return
	}

	Success(c, "ok", gin.H{
		"task_id":  task.TaskID,
		"status":   task.Status,
		"progress": task.Progress,
		"message":  task.Message,
	})
}

func (h *TaskHandler) Download(c *gin.Context) {
	taskID := c.Query("task_id")
	if taskID == "" {
		Fail(c, "task_id 不能为空")
		return
	}

	task, ok := service.GetTask(taskID)
	if !ok {
		Fail(c, "任务不存在或已过期")
		return
	}

	if task.Status != service.StatusDone {
		Fail(c, "任务尚未完成")
		return
	}

	if task.OutputPath == "" || !fileExists(task.OutputPath) {
		Fail(c, "输出文件不存在")
		return
	}

	// 标记为已下载，让清理协程在延迟后删除
	service.MarkDownloaded(taskID)

	c.FileAttachment(task.OutputPath, "clean_video.mp4")
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}