Yi.Admin/Yi.Ai.Vue3/src/pages/chat/image/components/ImageGenerator.vue

<script setup lang="ts">
import type { UploadFile, UploadUserFile } from 'element-plus';
import type { ImageModel, TaskStatusResponse } from '@/api/aiImage/types';
import {
  CircleCloseFilled,
  CopyDocument,
  Delete,
  Download,
  MagicStick,
  Picture as PictureIcon,
  Plus,
  Refresh,
  ZoomIn,
} from '@element-plus/icons-vue';
import { useClipboard } from '@vueuse/core';
import { ElMessage } from 'element-plus';
import { computed, onMounted, onUnmounted, ref, watch } from 'vue';
import { getSelectableTokenInfo } from '@/api';
import { generateImage, getImageModels, getTaskStatus } from '@/api/aiImage';

const props = defineProps({
  isActive: {
    type: Boolean,
    default: true,
  },
});

const emit = defineEmits(['task-created']);

const { copy } = useClipboard();

// State
const tokenOptions = ref<any[]>([]);
const selectedTokenId = ref('');
const tokenLoading = ref(false);

const modelOptions = ref<ImageModel[]>([]);
const selectedModelId = ref('');
const modelLoading = ref(false);

const prompt = ref('');
const fileList = ref<UploadUserFile[]>([]);
const compressImage = ref(true);

const generating = ref(false);
const currentTaskId = ref('');
const currentTask = ref<TaskStatusResponse | null>(null);
const showViewer = ref(false);
let pollTimer: any = null;
let debounceTimer: any = null;

const canGenerate = computed(() => {
  return selectedModelId.value && prompt.value && !generating.value;
});

// Watch isActive to manage polling
watch(() => props.isActive, (active) => {
  if (active) {
    // Resume polling if we have a processing task
    if (currentTaskId.value && currentTask.value?.taskStatus === 'Processing') {
      startPolling(currentTaskId.value);
    }
  }
  else {
    stopPolling();
  }
});

// Methods
async function fetchTokens() {
  tokenLoading.value = true;
  try {
    const res = await getSelectableTokenInfo();
    // Handle potential wrapper
    const data = Array.isArray(res) ? res : (res as any).data || [];

    // Add Default Option
    tokenOptions.value = [
      { tokenId: '', name: '默认 (Default)', isDisabled: false },
      ...data,
    ];

    // Default select "Default" if available, otherwise first available
    if (!selectedTokenId.value) {
      selectedTokenId.value = '';
    }
  }
  catch (e) {
    console.error(e);
  }
  finally {
    tokenLoading.value = false;
  }
}

async function fetchModels() {
  modelLoading.value = true;
  try {
    const res = await getImageModels();
    // Handle potential wrapper
    const data = Array.isArray(res) ? res : (res as any).data || [];
    modelOptions.value = data;

    // Default select first
    if (modelOptions.value.length > 0 && !selectedModelId.value) {
      selectedModelId.value = modelOptions.value[0].modelId;
    }
  }
  catch (e) {
    console.error(e);
  }
  finally {
    modelLoading.value = false;
  }
}

function handleFileChange(uploadFile: UploadFile) {
  const isLt5M = uploadFile.size! / 1024 / 1024 < 5;
  if (!isLt5M) {
    ElMessage.error('图片大小不能超过 5MB!');
    const index = fileList.value.indexOf(uploadFile);
    if (index !== -1)
      fileList.value.splice(index, 1);
  }
}

function handleRemove(file: UploadFile) {
  const index = fileList.value.indexOf(file);
  if (index !== -1)
    fileList.value.splice(index, 1);
}

function fileToBase64(file: File): Promise<string> {
  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.readAsDataURL(file);
    reader.onload = () => resolve(reader.result as string);
    reader.onerror = error => reject(error);
  });
}

async function handleGenerate() {
  if (!canGenerate.value)
    return;

  // Clear existing debounce timer
  if (debounceTimer) {
    clearTimeout(debounceTimer);
  }

  // Set debounce timer
  debounceTimer = setTimeout(async () => {
    await executeGenerate();
  }, 300);
}

async function executeGenerate() {
  generating.value = true;
  // Reset current task display immediately
  currentTask.value = {
    id: '',
    prompt: prompt.value,
    taskStatus: 'Processing',
    publishStatus: 'Unpublished',
    categories: [],
    creationTime: new Date().toISOString(),
  };

  try {
    const base64Images: string[] = [];
    for (const fileItem of fileList.value) {
      if (fileItem.raw) {
        const base64 = await fileToBase64(fileItem.raw);
        base64Images.push(base64);
      }
    }

    const res = await generateImage({
      tokenId: selectedTokenId.value || undefined, // Send undefined if empty
      prompt: prompt.value,
      modelId: selectedModelId.value,
      referenceImagesPrefixBase64: base64Images,
    });

    // Robust ID extraction: Handle string, object wrapper, or direct object
    let taskId = '';
    if (typeof res === 'string') {
      taskId = res;
    }
    else if (typeof res === 'object' && res !== null) {
      // Check for data property which might contain the ID string or object
      const data = (res as any).data;
      if (typeof data === 'string') {
        taskId = data;
      }
      else if (typeof data === 'object' && data !== null) {
        taskId = data.id || data.taskId;
      }
      else {
        // Fallback to direct properties
        taskId = (res as any).id || (res as any).taskId || (res as any).result;
      }
    }

    if (!taskId) {
      console.error('Task ID not found in response:', res);
      throw new Error('Invalid Task ID');
    }

    currentTaskId.value = taskId;
    startPolling(taskId);
    emit('task-created');
  }
  catch (e) {
    console.error(e);
    ElMessage.error('生成任务创建失败');
    currentTask.value = null;
  }
  finally {
    generating.value = false; // Allow new tasks immediately
  }
}

function startPolling(taskId: string) {
  if (pollTimer)
    clearInterval(pollTimer);

  // Initial fetch
  pollStatus(taskId);

  pollTimer = setInterval(() => {
    pollStatus(taskId);
  }, 3000);
}

async function pollStatus(taskId: string) {
  // Double check active status before polling (though timer should be cleared)
  if (!props.isActive) {
    stopPolling();
    return;
  }

  try {
    const res = await getTaskStatus(taskId);
    // Handle response structure if needed
    const taskData = (res as any).data || res;

    // Only update if it matches the current task we are watching
    // This prevents race conditions if user starts a new task
    if (currentTaskId.value === taskId) {
      currentTask.value = taskData;

      // Case-insensitive check just in case
      const status = taskData.taskStatus;

      if (status === 'Success' || status === 'Fail') {
        stopPolling();
        if (status === 'Success') {
          ElMessage.success('图片生成成功');
        }
        else {
          ElMessage.error(taskData.errorInfo || '图片生成失败');
        }
      }
    }
    else {
      // If task ID changed, stop polling this old task
      // Actually, we should probably just let the new task's poller handle it
      // But since we use a single pollTimer variable, starting a new task clears the old timer.
      // So this check is mostly for the initial async call returning after new task started.
    }
  }
  catch (e) {
    console.error(e);
    // Don't stop polling on transient network errors, but maybe log it
  }
}

function stopPolling() {
  if (pollTimer) {
    clearInterval(pollTimer);
    pollTimer = null;
  }
}

function clearPrompt() {
  prompt.value = '';
}

async function downloadImage() {
  if (currentTask.value?.storeUrl) {
    try {
      const response = await fetch(currentTask.value.storeUrl);
      const blob = await response.blob();
      const url = window.URL.createObjectURL(blob);
      const link = document.createElement('a');
      link.href = url;
      link.download = `generated-image-${currentTask.value.id}.png`;
      document.body.appendChild(link);
      link.click();
      document.body.removeChild(link);
      window.URL.revokeObjectURL(url);
    }
    catch (e) {
      console.error('Download failed', e);
      // Fallback
      window.open(currentTask.value.storeUrl, '_blank');
    }
  }
}

async function copyError(text: string) {
  await copy(text);
  ElMessage.success('错误信息已复制');
}

// Exposed methods for external control
function setPrompt(text: string) {
  prompt.value = text;
}

// Helper to load image from URL and convert to File object
async function addReferenceImage(url: string) {
  if (fileList.value.length >= 2) {
    ElMessage.warning('最多只能上传2张参考图');
    return;
  }

  try {
    const response = await fetch(url);
    const blob = await response.blob();
    const filename = url.split('/').pop() || 'reference.png';
    const file = new File([blob], filename, { type: blob.type });

    const uploadFile: UploadUserFile = {
      name: filename,
      url,
      raw: file,
      uid: Date.now(),
      status: 'ready',
    };

    fileList.value.push(uploadFile);
  }
  catch (e) {
    console.error('Failed to load reference image', e);
    ElMessage.error('无法加载参考图');
  }
}

defineExpose({
  setPrompt,
  addReferenceImage,
});

onMounted(() => {
  fetchTokens();
  fetchModels();
});

onUnmounted(() => {
  stopPolling();
});
</script>

<template>
  <div class="flex flex-col h-full md:flex-row gap-4 md:gap-6 p-3 md:p-4 bg-white rounded-lg shadow-sm">
    <!-- Left Config Panel -->
    <div class="w-full md:w-[400px] flex flex-col gap-4 md:gap-6 overflow-y-auto pr-0 md:pr-2 custom-scrollbar">
      <div class="space-y-3 md:space-y-4">
        <h2 class="text-base md:text-lg font-bold text-gray-800 flex items-center gap-2">
          <el-icon><MagicStick /></el-icon>
          配置
        </h2>

        <!-- 操作说明 -->
        <div class="px-3 py-2 bg-green-50 rounded-lg text-xs text-gray-600 leading-relaxed">
          <p class="mb-1"><strong>📌 生成说明：</strong></p>
          <p class="mb-1">• 生成的图片将自动保存到"我的图库"，仅个人可见</p>
          <p>• 发布后图片将在"图片广场"公开展示，所有用户可见</p>
        </div>

        <el-form label-position="top" class="space-y-1">
          <!-- Token -->
          <el-form-item label="API密钥 (可选)">
            <el-select
              v-model="selectedTokenId"
              placeholder="请选择API密钥"
              class="w-full"
              :loading="tokenLoading"
              clearable
            >
              <el-option
                v-for="token in tokenOptions"
                :key="token.tokenId"
                :label="token.name"
                :value="token.tokenId"
                :disabled="token.isDisabled"
              />
            </el-select>
          </el-form-item>

          <!-- Model -->
          <el-form-item label="模型" required>
            <el-select
              v-model="selectedModelId"
              placeholder="请选择模型"
              class="w-full"
              :loading="modelLoading"
            >
              <el-option
                v-for="model in modelOptions"
                :key="model.modelId"
                :label="model.modelName"
                :value="model.modelId"
              >
                <div class="flex flex-col py-1 max-w-[350px]">
                  <span class="font-medium truncate">{{ model.modelName }}</span>
                  <span class="text-xs text-gray-400 truncate" :title="model.modelDescribe">{{ model.modelDescribe }}</span>
                </div>
              </el-option>
            </el-select>
            <div class="mt-2 text-xs text-gray-500">
              按此收费，
              <router-link to="/model-library" class="text-primary hover:underline">
                前往模型库查看详情
              </router-link>
            </div>
          </el-form-item>

          <!-- Prompt -->
          <el-form-item label="提示词" required class="prompt-from">
            <template #label>
              <div class="flex justify-between items-center w-full">
                <span>提示词</span>
                <el-button link type="primary" size="small" @click="clearPrompt">
                  <el-icon class="mr-1">
                    <Delete />
                  </el-icon>清空
                </el-button>
              </div>
            </template>
            <el-input
              v-model="prompt"
              type="textarea"
              :autosize="{ minRows: 4, maxRows: 15 }"
              placeholder="描述你想要生成的画面，例如：一只在太空中飞行的赛博朋克风格的猫..."
              maxlength="2000"
              show-word-limit
              class="custom-textarea"
              resize="vertical"
            />
          </el-form-item>

          <!-- Reference Image -->
          <el-form-item label="参考图 (可选)">
            <div class="w-full bg-gray-50 p-4 rounded-lg border border-dashed border-gray-300 hover:border-blue-400 transition-colors">
              <el-upload
                v-model:file-list="fileList"
                action="#"
                list-type="picture-card"
                :auto-upload="false"
                :limit="2"
                :on-change="handleFileChange"
                :on-remove="handleRemove"
                accept=".jpg,.jpeg,.png,.bmp,.webp"
                :class="{ 'hide-upload-btn': fileList.length >= 2 }"
              >
                <div class="flex flex-col items-center justify-center text-gray-400">
                  <el-icon class="text-2xl mb-2">
                    <Plus />
                  </el-icon>
                  <span class="text-xs">点击上传</span>
                </div>
              </el-upload>
              <div class="text-xs text-gray-400 mt-2 flex justify-between items-center flex-wrap gap-2">
                <span>最多2张，< 5MB (支持 JPG/PNG/WEBP)</span>
                <el-checkbox v-model="compressImage" label="压缩图片" size="small" />
              </div>
            </div>
          </el-form-item>
        </el-form>
      </div>

      <div class="mt-auto pt-1">
        <el-button
          type="primary"
          class="w-full h-12 text-lg shadow-lg shadow-blue-500/30 transition-all hover:shadow-blue-500/50"
          :loading="generating"
          :disabled="!canGenerate"
          round
          @click="handleGenerate"
        >
          {{ generating ? '生成中...' : '开始生成' }}
        </el-button>
      </div>
    </div>

    <!-- Right Result Panel -->
    <div class="flex-1 bg-gray-100 rounded-xl overflow-hidden relative flex flex-col h-full">
      <div v-if="currentTask" class="flex-1 flex flex-col relative h-full overflow-hidden">
        <!-- Image Display -->
        <div class="flex-1 flex items-center justify-center p-4 bg-checkboard overflow-hidden relative min-h-0">
          <div v-if="currentTask.taskStatus === 'Success' && currentTask.storeUrl" class="relative group w-full h-full flex items-center justify-center">
            <el-image
              ref="previewImageRef"
              :src="currentTask.storeUrl"
              fit="contain"
              class="w-full h-full"
              :style="{ maxHeight: '100%', maxWidth: '100%' }"
              :preview-src-list="[currentTask.storeUrl]"
              :initial-index="0"
              :preview-teleported="true"
            />
            <!-- Hover Actions -->
            <div class="absolute bottom-4 right-4 opacity-0 group-hover:opacity-100 transition-opacity flex gap-2 z-10">
              <el-button circle type="primary" :icon="ZoomIn" @click="showViewer = true" />
              <el-button circle type="primary" :icon="Download" @click="downloadImage" />
              <el-button circle type="success" :icon="Refresh" title="重新生成" @click="handleGenerate" />
            </div>

            <el-image-viewer
              v-if="showViewer"

              fit="contain"
              :url-list="[currentTask.storeUrl]"
              @close="showViewer = false"
            />
          </div>

          <!-- Processing State -->
          <div v-else-if="currentTask.taskStatus === 'Processing'" class="text-center">
            <div class="loader mb-6" />
            <p class="text-gray-600 font-medium text-lg animate-pulse">
              正在绘制您的想象...
            </p>
            <p class="text-gray-400 text-sm mt-2">
              请稍候，这可能需要一小会~
            </p>

            <p class="text-gray-400 text-sm mt-2">
              您可以离开或者继续创建下一个任务，稍后到“我的图库”中查看任务进度
            </p>
          </div>

          <!-- Fail State -->
          <div v-else-if="currentTask.taskStatus === 'Fail'" class="text-center text-red-500 w-full max-w-lg">
            <el-icon class="text-6xl mb-4">
              <CircleCloseFilled />
            </el-icon>
            <p class="text-lg font-medium mb-2">
              生成失败
            </p>
            <div class="bg-red-50 p-4 rounded-lg border border-red-100 text-sm text-left relative group/error">
              <div class="max-h-32 overflow-y-auto custom-scrollbar break-words pr-6">
                {{ currentTask.errorInfo || '请检查提示词或稍后重试' }}
              </div>
              <el-button
                class="absolute top-2 right-2 opacity-0 group-hover/error:opacity-100 transition-opacity"
                size="small"
                circle
                :icon="CopyDocument"
                title="复制错误信息"
                @click.stop="copyError(currentTask.errorInfo || '')"
              />
            </div>
            <el-button class="mt-6" icon="Refresh" @click="handleGenerate">
              重试
            </el-button>
          </div>
        </div>

        <!-- Prompt Display (Bottom) -->
        <div class="bg-white p-4 border-t border-gray-200 shrink-0 max-h-40 overflow-y-auto">
          <p class="text-xs font-bold text-gray-500 uppercase tracking-wider mb-2">
            当前提示词
          </p>
          <p class="text-gray-800 text-sm leading-relaxed whitespace-pre-wrap">
            {{ currentTask.prompt }}
          </p>
        </div>
      </div>

      <!-- Empty State -->
      <div v-else class="flex-1 flex flex-col items-center justify-center text-gray-400 bg-gray-50/50 h-full">
        <div class="w-32 h-32 bg-gray-200 rounded-full flex items-center justify-center mb-6">
          <el-icon class="text-5xl text-gray-400">
            <PictureIcon />
          </el-icon>
        </div>
        <h3 class="text-xl font-semibold text-gray-600 mb-2">
          准备好开始了吗？
        </h3>
        <p class="text-gray-500 max-w-md text-center">
          在左侧配置您的创意参数，点击生成按钮，见证AI的奇迹。
        </p>
      </div>
    </div>
  </div>
</template>

<style scoped>
.custom-scrollbar::-webkit-scrollbar {
  width: 6px;
}
.custom-scrollbar::-webkit-scrollbar-thumb {
  background-color: #e5e7eb;
  border-radius: 3px;
}
.custom-scrollbar::-webkit-scrollbar-track {
  background-color: transparent;
}
:deep(.prompt-from .el-form-item__label){
  display: flex;
}
.bg-checkboard {
  background-image:
    linear-gradient(45deg, #f0f0f0 25%, transparent 25%),
    linear-gradient(-45deg, #f0f0f0 25%, transparent 25%),
    linear-gradient(45deg, transparent 75%, #f0f0f0 75%),
    linear-gradient(-45deg, transparent 75%, #f0f0f0 75%);
  background-size: 20px 20px;
  background-position: 0 0, 0 10px, 10px -10px, -10px 0px;
}

/* Loader Animation */
.loader {
  width: 48px;
  height: 48px;
  border: 5px solid #e5e7eb;
  border-bottom-color: #3b82f6;
  border-radius: 50%;
  display: inline-block;
  box-sizing: border-box;
  animation: rotation 1s linear infinite;
}

@keyframes rotation {
  0% {
    transform: rotate(0deg);
  }
  100% {
    transform: rotate(360deg);
  }
}

:deep(.hide-upload-btn .el-upload--picture-card) {
  display: none;
}
</style>