furniture-3d-system-Wu/furniture-3d-system/src/app/ai-render-workbench/sceneGeneration.ts

'use client'

import type { CameraController } from '@/lib/3d/camera-controller'
import {
  createAssistedAssemblyAiRender,
  waitForAssistedAssemblyAiRender,
  type CreateAssistedAssemblyAiRenderPayload,
} from '@/lib/open-product/platform-api'
import {
  buildLockedScenePrompt,
  CURVE_DETAIL_ENHANCEMENT_NEGATIVE_PROMPT,
  WHITE_BACKGROUND_COMPONENT_TYPE_LOCK_NEGATIVE_PROMPT,
  WHITE_BACKGROUND_MATERIAL_SCOPE_LOCK_NEGATIVE_PROMPT,
} from './promptBuilders'
import { mergeNegativePrompt, buildTaskStartedMessage } from './workbenchHelpers'
import type {
  AiRenderCameraViewSnapshot,
  AiRenderSku,
  AiRenderSpu,
  AiRenderTaskState,
  LocalRenderRecord,
} from './types'
import { buildAiRenderAssetSnapshot, buildAiRenderWorkspaceSnapshot } from './utils'

type SceneGenerationParams = {
  canGenerate: boolean
  generating: boolean
  selectedSku: AiRenderSku | null
  selectedSpu: AiRenderSpu | null
  prompt: string
  negativePrompt: string
  aspectRatio: string
  imageCount: number
  imageSize: CreateAssistedAssemblyAiRenderPayload['imageSize']
  modelNodeNames: string[]
  getCurrentCameraView: () => AiRenderCameraViewSnapshot | null
  getController: () => CameraController | null
  captureCurrentReferenceImage: () => string | null
  captureReferenceAtView: (
    view: AiRenderCameraViewSnapshot,
    options?: {
      controller?: CameraController | null
      restoreAfterCapture?: boolean
    },
  ) => Promise<string | null>
  startRenderTask: (params: {
    mode: LocalRenderRecord['mode']
    skuId: string
    skuName: string
    imageCount: number
  }) => AiRenderTaskState | null
  appendServerRenderRecords: (render: any, skuFallback?: AiRenderSku | null) => number
  setMessageForSku: (skuId: string, nextMessage: string | null) => void
  settleRenderTask: (taskId: string, nextStatus: 'succeeded' | 'failed', stage: string) => void
  setMessage: (message: string | null) => void
}

export async function runSceneGeneration(params: SceneGenerationParams): Promise<void> {
  const {
    canGenerate,
    generating,
    selectedSku,
    selectedSpu,
    prompt,
    negativePrompt,
    aspectRatio,
    imageCount,
    imageSize,
    modelNodeNames,
    getCurrentCameraView,
    getController,
    captureCurrentReferenceImage,
    captureReferenceAtView,
    startRenderTask,
    appendServerRenderRecords,
    setMessageForSku,
    settleRenderTask,
    setMessage,
  } = params

  if (!canGenerate || generating) return

  const taskSku = selectedSku
  const taskSpu = selectedSpu
  const controller = getController()
  if (!taskSku || !taskSku.modelUrl) {
    setMessage('请先选择并配置一个 SKU 模型文件。')
    return
  }

  const cameraView = getCurrentCameraView()
  const referenceImageDataUrl =
    cameraView && controller
      ? await captureReferenceAtView(cameraView, {
          controller,
          restoreAfterCapture: true,
        })
      : captureCurrentReferenceImage()
  if (!referenceImageDataUrl) {
    setMessage('当前 3D 预览还没准备好，无法作为 AI 参考图。')
    return
  }

  const normalizedPrompt = prompt.trim()
  if (!normalizedPrompt) {
    setMessage('请先填写场景效果图提示词。')
    return
  }

  const taskImageCount = imageCount
  const taskAspectRatio = aspectRatio
  const taskImageSize = imageSize
  const taskNegativePrompt = negativePrompt
  const taskNodeNames = [...modelNodeNames]
  const lockedScenePrompt = buildLockedScenePrompt({
    spu: taskSpu,
    sku: taskSku,
    userPrompt: normalizedPrompt,
    aspectRatio: taskAspectRatio,
  })
  const task = startRenderTask({
    mode: 'scene',
    skuId: taskSku.id,
    skuName: taskSku.name || '未命名 SKU',
    imageCount: taskImageCount,
  })
  if (!task) {
    setMessage('当前 SKU 已有生图任务在后台运行。')
    return
  }

  setMessage(buildTaskStartedMessage(task))

  try {
    const pending = await createAssistedAssemblyAiRender({
      prompt: lockedScenePrompt,
      negativePrompt: mergeNegativePrompt(
        taskNegativePrompt,
        WHITE_BACKGROUND_COMPONENT_TYPE_LOCK_NEGATIVE_PROMPT,
        WHITE_BACKGROUND_MATERIAL_SCOPE_LOCK_NEGATIVE_PROMPT,
        CURVE_DETAIL_ENHANCEMENT_NEGATIVE_PROMPT,
      ),
      imageCount: taskImageCount,
      aspectRatio: taskAspectRatio,
      imageSize: taskImageSize,
      mode: 'scene',
      intakeSkuItemId: taskSku.id,
      referenceImageDataUrl,
      workspaceSnapshot: buildAiRenderWorkspaceSnapshot({
        spu: taskSpu,
        sku: taskSku,
        prompt: lockedScenePrompt,
        mode: 'scene',
        cameraView,
      }),
      assetSnapshot: buildAiRenderAssetSnapshot({
        spu: taskSpu,
        sku: taskSku,
        nodeNames: taskNodeNames,
      }),
    }, { onUncertain: (message) => setMessageForSku(taskSku.id, message) })
    appendServerRenderRecords(pending, taskSku)
    if (pending.submissionRecovery) {
      setMessageForSku(taskSku.id, '已找回上次提交的任务，请先核对原记录。本次没有新建生成任务。')
      settleRenderTask(task.id, 'succeeded', '已恢复原任务，未新建任务')
      return
    }
    const render = await waitForAssistedAssemblyAiRender(pending)
    const savedCount = appendServerRenderRecords(render, taskSku)
    if (render.status === 'SUCCEEDED' && savedCount > 0) {
      setMessageForSku(
        taskSku.id,
        render.errorMessage || (savedCount > 1 ? `场景效果图已生成，共 ${savedCount} 张。` : '场景效果图已生成。'),
      )
      settleRenderTask(task.id, 'succeeded', `已完成，生成 ${savedCount} 张图片`)
      return
    }

    setMessageForSku(taskSku.id, render.errorMessage || 'AI 生图请求已返回，但没有可用结果图。')
    settleRenderTask(task.id, 'failed', render.status === 'PENDING' ? '任务待核对，请勿重复提交' : '生成结束，但没有返回可用图片')
  } catch (error) {
    setMessageForSku(taskSku.id, error instanceof Error ? error.message : 'AI 生图失败')
    settleRenderTask(task.id, 'failed', '生成失败，请稍后重试')
  }
}