gen2d/frontend/src/stores/generation.ts

import { create } from 'zustand'
import type { Asset, GenerateRequest, PipelineStage } from '../api/types'
import { submitGenerate, getTask, getAssets } from '../api/generate'

type Status = 'idle' | 'submitting' | 'running' | 'completed' | 'failed'

interface GenerationState {
  taskId: string | null
  projectId: string | null
  stage: PipelineStage | null
  progress: number
  status: Status
  statusText: string
  retryCount: number
  rejectReason: string | null
  assets: Asset[]
  error: string | null
  submit: (req: GenerateRequest) => Promise<void>
  reset: () => void
}

let pollTimer: ReturnType<typeof setInterval> | null = null

function stopPolling() {
  if (pollTimer) {
    clearInterval(pollTimer)
    pollTimer = null
  }
}

export const useGenerationStore = create<GenerationState>((set) => ({
  taskId: null,
  projectId: null,
  stage: null,
  progress: 0,
  status: 'idle',
  statusText: '',
  retryCount: 0,
  rejectReason: null,
  assets: [],
  error: null,

  submit: async (req) => {
    stopPolling()
    set({ status: 'submitting', error: null, statusText: '提交中...', stage: null })
    try {
      const { taskId } = await submitGenerate(req)

      set({
        taskId,
        projectId: req.projectId,
        status: 'running',
        progress: 5,
        stage: 'prompt_builder',
        statusText: '任务已提交，等待生成...',
      })

      pollTimer = setInterval(async () => {
        try {
          const task = await getTask(taskId)

          set((s) => ({
            stage: task.stage ?? s.stage,
            progress: task.progress ?? s.progress,
            retryCount: task.retryCount ?? s.retryCount,
            statusText:
              task.status === 'running'
                ? stageLabel(task.stage)
                : task.status === 'pending'
                  ? '排队中...'
                  : task.status,
          }))

          if (task.status === 'completed') {
            stopPolling()
            set({ progress: 90, statusText: '获取结果...' })

            const assets = await getAssets(taskId)
            set({
              status: 'completed',
              progress: 100,
              stage: 'format_adapter',
              statusText: '生成完成',
              assets,
            })
          } else if (task.status === 'failed') {
            stopPolling()
            set({
              status: 'failed',
              error: task.error || '生成失败',
              statusText: '生成失败',
            })
          }
        } catch {
          // 网络错误不中断轮询
        }
      }, 1500)
    } catch (err) {
      stopPolling()
      set({ status: 'failed', error: (err as Error).message, statusText: '提交失败' })
    }
  },

  reset: () => {
    stopPolling()
    set({
      taskId: null,
      projectId: null,
      stage: null,
      progress: 0,
      status: 'idle',
      statusText: '',
      retryCount: 0,
      rejectReason: null,
      assets: [],
      error: null,
    })
  },
}))

function stageLabel(stage?: string): string {
  switch (stage) {
    case 'prompt_builder': return '优化提示词...'
    case 'asset_generator': return '生成素材中...'
    case 'quality_supervisor': return '质检中...'
    case 'format_adapter': return '格式转换中...'
    default: return '生成中...'
  }
}