gen2d/backend/internal/handler/generate.go

package handler

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"gen2d/internal/model"
	"gen2d/internal/service"

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

// GenerateRequest 素材生成请求。
type GenerateRequest struct {
	ProjectID    string            `json:"projectId"`
	Prompt       string            `json:"prompt"`
	AssetType    string            `json:"assetType" binding:"required"`
	Tags         []string          `json:"tags"`
	UserNote     string            `json:"userNote"`
	ProjectStyle map[string]string `json:"projectStyle"`
	TaskStyle    map[string]string `json:"taskStyle"`
	Resolution   int               `json:"resolution"`
	Directions   int               `json:"directions"`
	FramesPerDir int               `json:"framesPerDir"`
	Format       string            `json:"format"`
}

// GenerateResponse 素材生成响应体。
type GenerateResponse struct {
	TaskID string `json:"taskId"`
}

// AssetResponse 单个素材响应。
type AssetResponse struct {
	URL    string `json:"url"`
	Format string `json:"format"`
}

// TaskResponse 任务查询响应。
type TaskResponse struct {
	ID         string `json:"id"`
	ProjectID  string `json:"projectId"`
	Prompt     string `json:"prompt"`
	AssetType  string `json:"assetType"`
	Status     string `json:"status"`
	Stage      string `json:"stage,omitempty"`
	Progress   int    `json:"progress"`
	RetryCount int    `json:"retryCount"`
	Error      string `json:"error,omitempty"`
	CreatedAt  string `json:"createdAt"`
}

// AssetsResponse 素材列表响应。
type AssetsResponse struct {
	Assets   []AssetResponse        `json:"assets"`
	Metadata service.AssetMetadata  `json:"metadata"`
}

// taskRecord 内存中的任务记录。
type taskRecord struct {
	task     TaskResponse
	assets   []AssetResponse
	metadata service.AssetMetadata
}

var (
	taskStore = sync.Map{} // taskID → *taskRecord
)

// Generate 素材生成接口（异步）。
// 立即返回 taskId，后台执行管线，前端通过 GET /tasks/:taskId 轮询进度。
func Generate(c *gin.Context) {
	var req GenerateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "参数错误: "+err.Error()))
		return
	}

	projectID := req.ProjectID
	if projectID == "" {
		projectID = "default"
	}
	taskID := fmt.Sprintf("task-%d", time.Now().UnixMilli())
	createdAt := time.Now().Format(time.RFC3339)

	// 存入 pending 状态
	taskStore.Store(taskID, &taskRecord{
		task: TaskResponse{
			ID:        taskID,
			ProjectID: projectID,
			Prompt:    req.Prompt,
			AssetType: req.AssetType,
			Status:    "pending",
			Progress:  0,
			CreatedAt: createdAt,
		},
	})

	// 返回 taskId
	c.JSON(http.StatusOK, model.OK(GenerateResponse{TaskID: taskID}))

	// 后台执行管线
	go runPipelineBg(projectID, taskID, req)
}

// runPipelineBg 后台执行生成管线，更新任务状态。
func runPipelineBg(projectID, taskID string, req GenerateRequest) {
	// 注入进度上报回调
	ctx := service.WithProgressReporter(context.Background(), func(stage string, progress int) {
		updateTaskProgress(taskID, "running", stage, progress)
	})

	updateTaskProgress(taskID, "running", "prompt_builder", 5)

	in := service.PipelineInput{
		ProjectID:    projectID,
		TaskID:       taskID,
		Prompt:       req.Prompt,
		AssetType:    req.AssetType,
		Tags:         req.Tags,
		UserNote:     req.UserNote,
		ProjectStyle: req.ProjectStyle,
		TaskStyle:    req.TaskStyle,
		Params: service.AssetParams{
			Resolution: req.Resolution,
			Frames: service.FrameParams{
				Directions:         req.Directions,
				FramesPerDirection: req.FramesPerDir,
			},
			Format: req.Format,
		},
	}

	output, err := service.RunPipeline(ctx, in)
	if err != nil {
		log.Printf("[generate] task %s failed: %v", taskID, err)
		updateFailed(taskID, err.Error())
		return
	}

	updateTaskProgress(taskID, "saving", "format_adapter", 90)

	// 保存图片到 ../generation/{projectId}/{taskId}/
	genDir := filepath.Join("..", "generation", projectID, taskID)
	if err := os.MkdirAll(genDir, 0755); err != nil {
		updateFailed(taskID, "创建输出目录失败: "+err.Error())
		return
	}

	assets := make([]AssetResponse, len(output.Assets))
	for i, a := range output.Assets {
		filename := fmt.Sprintf("%d.%s", i, a.Format)
		filePath := filepath.Join(genDir, filename)
		if err := os.WriteFile(filePath, a.Data, 0644); err != nil {
			updateFailed(taskID, "保存图片失败: "+err.Error())
			return
		}
		assets[i] = AssetResponse{
			URL:    fmt.Sprintf("/generation/%s/%s/%s", projectID, taskID, filename),
			Format: a.Format,
		}
	}

	// 更新为完成状态
	taskStore.Store(taskID, &taskRecord{
		task: TaskResponse{
			ID:        taskID,
			ProjectID: projectID,
			Prompt:    req.Prompt,
			AssetType: req.AssetType,
			Status:    "completed",
			Progress:  100,
			CreatedAt: time.Now().Format(time.RFC3339),
		},
		assets:   assets,
		metadata: output.Metadata,
	})

	log.Printf("[generate] task %s completed, %d assets", taskID, len(assets))
}

func updateTaskProgress(taskID, status, stage string, progress int) {
	rec, ok := taskStore.Load(taskID)
	if !ok {
		return
	}
	r := rec.(*taskRecord)
	r.task.Status = status
	r.task.Stage = stage
	r.task.Progress = progress
	taskStore.Store(taskID, r)
}

func updateFailed(taskID, errMsg string) {
	rec, ok := taskStore.Load(taskID)
	if !ok {
		return
	}
	r := rec.(*taskRecord)
	r.task.Status = "failed"
	r.task.Error = errMsg
	taskStore.Store(taskID, r)
}

// GetTask 查询任务信息。
func GetTask(c *gin.Context) {
	taskID := c.Param("taskId")
	rec, ok := taskStore.Load(taskID)
	if !ok {
		c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在"))
		return
	}
	r := rec.(*taskRecord)
	c.JSON(http.StatusOK, model.OK(r.task))
}

// GetAssets 查询任务素材列表。
func GetAssets(c *gin.Context) {
	taskID := c.Param("taskId")
	rec, ok := taskStore.Load(taskID)
	if !ok {
		c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在"))
		return
	}
	r := rec.(*taskRecord)
	if r.task.Status != "completed" {
		c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "任务尚未完成，当前状态: "+r.task.Status))
		return
	}
	c.JSON(http.StatusOK, model.OK(AssetsResponse{
		Assets:   r.assets,
		Metadata: r.metadata,
	}))
}