AI_Painting_V2.0/src/utils/taskPolling.js

import { ElNotification } from 'element-plus'
import { h } from 'vue'
import { requestCreateTask, requestTaskStatus } from '@/apis/display'
import { useDisplayStore, useUserStore } from '@/stores'
import { userError } from '@/utils/tokenError'

export function getChargeType(chargeType) {
  switch (chargeType) {
    case 'Painting':
      return 1
    case 'Video':
      return 4
    case 'Music':
      return 5
    default:
      return 2
  }
}

export function websocketError(code, msg) {
  let message
  switch (code) {
    case 1006:
      message = '用户身份验证失败'
      userError()
      break
    case 4401:
      message = msg
      break
    case 4402:
      message = JSON.parse(msg)
      break
    case 4403:
      message = msg
      break
    default:
      message = '连接异常，请稍后重试'
  }

  ElNotification({
    title: '生成失败',
    message: h('i', { style: 'color: teal' }, message),
    type: 'error',
    duration: 6000
  })
}

export function websocketSuccess() {
  ElNotification({
    title: '生成成功',
    message: h('div', [
      h('div', { style: 'font-weight: bold; color: teal;' }, '生成成功！'),
      h('br'),
      h('div', { style: 'color: orange; margin-top: 5px;' }, '内测状态，请及时下载生成的文件，云端储存与历史记录保留24小时！')
    ]),
    type: 'success',
    duration: 6000
  })
}

// 当前活跃的轮询定时器集合，用于页面卸载时清理
const activePollIntervals = new Set()

export async function generate(data, generateData) {
  const useDisplay = useDisplayStore()
  let taskId = null
  let pollInterval = null

  if (!data.modelId) {
    ElNotification({
      title: '生成失败',
      message: h('i', { style: 'color: teal' }, '未找到模型ID，请联系管理员配置'),
      type: 'error'
    })
    return
  }

  useDisplay.isSubGerenate = true

  // 从登录态获取 sessionId
  const sessionId = useUserStore().userInfo.sessionId
  if (!sessionId) {
    ElNotification({
      title: '生成失败',
      message: h('i', { style: 'color: teal' }, '用户身份已过期，请重新登录'),
      type: 'error'
    })
    useDisplay.isSubGerenate = false
    return
  }

  try {
    const body = data.body

    // 构造请求体
    const requestBody = {
      model_id: data.modelId,
      body,
      request: data.request
    }

    // POST 创建任务
    const createResult = await requestCreateTask(requestBody, sessionId)

    if (createResult.code !== 0) {
      ElNotification({
        title: '生成失败',
        message: h('i', { style: 'color: teal' }, createResult.message || '任务创建失败'),
        type: 'error'
      })
      useDisplay.isSubGerenate = false
      return
    }

    taskId = createResult.data.task_id

    // 在列表中插入"生成中"条目
    useDisplay.addGeneratingItem({
      taskId,
      type: data.type,
      generateData
    })
    setTimeout(() => {
      useDisplay.scrollToBottom()
    }, 100)

    // 轮询任务状态
    const pollTask = async () => {
      try {
        const pollResult = await requestTaskStatus(taskId)

        if (pollResult.code !== 0) return

        const taskData = pollResult.data

        if (taskData.status === 'completed') {
          clearInterval(pollInterval)
          activePollIntervals.delete(pollInterval)
          useDisplay.isSubGerenate = false

          // 提取结果 URL
          const urls = taskData.outputs?.map((img) => img.url) || []
          if (urls.length > 0) {
            useDisplay.updateItemToSuccess(taskId, urls)
            websocketSuccess()
          } else {
            websocketError(4403, '未获取到生成结果')
          }
        } else if (taskData.status === 'failed') {
          clearInterval(pollInterval)
          activePollIntervals.delete(pollInterval)
          useDisplay.isSubGerenate = false
          websocketError(4403, taskData.vendor_error || '生成失败')
        }
        // queued / processing 状态继续轮询
      } catch (error) {
        console.error('轮询任务状态失败:', error)
      }
    }

    // 每 20 秒轮询一次
    pollInterval = setInterval(pollTask, 20000)
    activePollIntervals.add(pollInterval)

    // 5 秒后先做第一次轮询
    setTimeout(pollTask, 5000)
  } catch (error) {
    console.error('创建任务失败:', error)
    useDisplay.isSubGerenate = false
    ElNotification({
      title: '生成通知',
      message: h('i', { style: 'color: teal' }, '生成失败，请检查参数后重新提交任务'),
      type: 'error'
    })
    if (pollInterval) {
      clearInterval(pollInterval)
      activePollIntervals.delete(pollInterval)
    }
  }
}

// 页面卸载时清理所有轮询
if (typeof window !== 'undefined') {
  window.addEventListener('beforeunload', () => {
    for (const interval of activePollIntervals) {
      clearInterval(interval)
    }
    activePollIntervals.clear()
  })
}