Atlas/web/src/composables/useImageStudio.ts

import { computed, onMounted, onUnmounted, ref } from 'vue'
import api from '@/api/client'
import { useUserMeta, type ModelInfo } from '@/composables/useUserMeta'
import type { TaskItem } from '@/composables/useInvoke'

const serviceType = 'image_generation'
const legacyImageServiceTypes = ['text_to_image', 'image_to_image'] as const
const terminalStatuses = new Set(['succeeded', 'failed', 'canceled'])

export function useImageStudio() {
  const { load, getService } = useUserMeta()
  const activeTasks = ref<TaskItem[]>([])
  const streams = new Map<number, EventSource>()

  const models = computed(() => (getService(serviceType)?.models ?? []) as ModelInfo[])

  const defaultModel = computed(() => {
    const id = getService(serviceType)?.default_model
    if (id && models.value.some((m) => m.id === id)) return id
    return models.value[0]?.id ?? ''
  })

  function upsertTask(task: TaskItem) {
    const idx = activeTasks.value.findIndex((t) => t.id === task.id)
    if (idx >= 0) activeTasks.value[idx] = task
    else activeTasks.value.push(task)
  }

  function closeStream(taskId: number) {
    const es = streams.get(taskId)
    if (es) {
      es.close()
      streams.delete(taskId)
    }
  }

  function watchTask(taskId: number) {
    if (streams.has(taskId)) return
    const es = new EventSource(`/api/v1/user/tasks/${taskId}/events`, { withCredentials: true })
    streams.set(taskId, es)

    es.addEventListener('task', (ev) => {
      const task = JSON.parse(ev.data) as TaskItem
      upsertTask(task)
      if (terminalStatuses.has(task.status)) closeStream(taskId)
    })
    es.addEventListener('done', () => closeStream(taskId))
    es.onerror = () => closeStream(taskId)
  }

  async function loadActiveTasks() {
    const types = [serviceType, ...legacyImageServiceTypes]
    const responses = await Promise.all(
      types.map((st) =>
        api.get<{ items: TaskItem[] }>('/user/tasks/active', { params: { service_type: st } }),
      ),
    )
    const items = responses.flatMap((res) => res.data.items)
    activeTasks.value = items
    for (const task of items) watchTask(task.id)
  }

  async function generate(params: Record<string, unknown>, model: string) {
    try {
      const res = await api.post('/user/generate', {
        service_type: serviceType,
        params,
        model,
      })
      if (res.data.task_id) {
        watchTask(res.data.task_id)
        await loadActiveTasks()
      }
      return res.data
    } catch (err: unknown) {
      const msg =
        (err as { response?: { data?: { error?: string } } })?.response?.data?.error ?? '生成请求失败'
      throw new Error(msg)
    }
  }

  onMounted(loadActiveTasks)
  onUnmounted(() => {
    streams.forEach((es) => es.close())
    streams.clear()
  })

  return {
    load,
    models,
    defaultModel,
    activeTasks,
    generate,
    loadActiveTasks,
  }
}

export function modelKindLabel(kind?: ModelInfo['kind']) {
  return kind === 'image_to_image' ? '图生图' : '文生图'
}