import { create } from 'zustand'
import type { Asset, PipelineProgress, PipelineStage } from '../api/types'
import { submitGenerate } from '../api/generate'
import { createMockWebSocket } from '../api/mock'

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

interface GenerationState {
  taskId: string | null
  stage: PipelineStage | null
  progress: number
  status: Status
  retryCount: number
  rejectReason: string | null
  assets: Asset[]
  error: string | null
  submit: (projectId: string, prompt: string, assetType: string) => Promise<void>
  handleProgress: (msg: PipelineProgress) => void
  reset: () => void
}

let cleanupWs: (() => void) | null = null

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

  submit: async (projectId, prompt, assetType) => {
    set({ status: 'submitting', error: null })
    try {
      const taskId = await submitGenerate(projectId, prompt, assetType)
      set({ taskId, status: 'running', progress: 0 })

      // 启动 mock WebSocket
      cleanupWs = createMockWebSocket(
        taskId,
        (msg) => get().handleProgress(msg),
        (assets) => {
          set({ status: 'completed', assets, progress: 100 })
        },
        (error) => {
          set({ status: 'failed', error })
        }
      )
    } catch (err) {
      set({ status: 'failed', error: (err as Error).message })
    }
  },

  handleProgress: (msg) => {
    set({
      stage: msg.stage,
      progress: msg.progress,
      retryCount: msg.retryCount ?? get().retryCount,
      rejectReason: msg.rejectReason ?? null,
    })
    if (msg.result?.assets) {
      set({ assets: msg.result.assets })
    }
  },

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