gen2d/frontend/src/stores/generation.ts

import { create } from 'zustand'
import type { Asset, GenerateRequest } 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
  progress: number
  status: Status
  statusText: string
  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, get) => ({
  taskId: null,
  projectId: null,
  progress: 0,
  status: 'idle',
  statusText: '',
  assets: [],
  error: null,

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

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

      // 开始轮询进度
      pollTimer = setInterval(async () => {
        try {
          const task = await getTask(taskId)

          set({
            progress: task.progress ?? get().progress,
            statusText:
              task.status === 'running'
                ? '生成中...'
                : task.status === 'pending'
                  ? '排队中...'
                  : task.status,
          })

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

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

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