AnimationManager/server/ai-queue.js

// AI生图队列管理器
const fs = require('fs');
const path = require('path');
const ReplaceCharacterHandler = require('./replace-character');

const AI_HISTORY_FILE = path.join(__dirname, 'ai-history.json');
const QUEUE_FILE = path.join(__dirname, 'ai-queue.json');

// 队列状态
let queue = [];
let isProcessing = false;

// 初始化：加载队列
function initQueue() {
  try {
    if (fs.existsSync(QUEUE_FILE)) {
      const data = fs.readFileSync(QUEUE_FILE, 'utf-8');
      queue = JSON.parse(data);
    }
  } catch (error) {
    console.error('[AIQueue] 加载队列失败:', error);
    queue = [];
  }
}

// 保存队列
function saveQueue() {
  try {
    fs.writeFileSync(QUEUE_FILE, JSON.stringify(queue, null, 2), 'utf-8');
  } catch (error) {
    console.error('[AIQueue] 保存队列失败:', error);
  }
}

// 加载AI历史
function loadAIHistory() {
  try {
    if (fs.existsSync(AI_HISTORY_FILE)) {
      const data = fs.readFileSync(AI_HISTORY_FILE, 'utf-8');
      return JSON.parse(data);
    }
  } catch (error) {
    console.error('[AIQueue] 加载历史失败:', error);
  }
  return {};
}

// 保存AI历史
function saveAIHistory(history) {
  try {
    fs.writeFileSync(AI_HISTORY_FILE, JSON.stringify(history, null, 2), 'utf-8');
  } catch (error) {
    console.error('[AIQueue] 保存历史失败:', error);
  }
}

// 添加任务到队列
function addToQueue(username, taskData) {
  const taskId = `ai_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  const task = {
    id: taskId,
    username: username.toLowerCase(),
    status: queue.length === 0 && !isProcessing ? 'rendering' : 'queued',
    createdAt: new Date().toISOString(),
    ...taskData
  };
  
  queue.push(task);
  saveQueue();
  
  // 保存原始图片预览（用于模糊显示）
  let previewUrl = null;
  if (taskData.image1) {
    previewUrl = savePreviewImage(username, taskId, taskData.image1);
  }
  
  // 添加到历史记录
  const history = loadAIHistory();
  if (!history[task.username]) {
    history[task.username] = [];
  }
  history[task.username].unshift({
    id: taskId,
    status: task.status,
    createdAt: task.createdAt,
    imageUrl: null,
    previewUrl: previewUrl // 原始图片预览
  });
  saveAIHistory(history);
  
  // 如果队列为空且没有正在处理的任务，立即开始处理
  if (queue.length === 1 && !isProcessing) {
    processQueue();
  }
  
  return taskId;
}

// 保存预览图片（原始texture的缩略图）
function savePreviewImage(username, taskId, imageBase64) {
  const usersDir = path.join(__dirname, 'users');
  const userDir = path.join(usersDir, username.toLowerCase());
  const aiDir = path.join(userDir, 'ai-images');
  
  // 确保目录存在
  if (!fs.existsSync(aiDir)) {
    fs.mkdirSync(aiDir, { recursive: true });
  }
  
  const imagePath = path.join(aiDir, `${taskId}_preview.png`);
  const imageBuffer = Buffer.from(imageBase64.replace(/^data:image\/\w+;base64,/, ''), 'base64');
  fs.writeFileSync(imagePath, imageBuffer);
  
  return `/api/ai/preview?username=${encodeURIComponent(username)}&id=${encodeURIComponent(taskId)}`;
}

// 处理队列
async function processQueue() {
  if (isProcessing || queue.length === 0) {
    return;
  }
  
  isProcessing = true;
  
  while (queue.length > 0) {
    const task = queue[0];
    
    // 更新状态为rendering
    if (task.status === 'queued') {
      task.status = 'rendering';
      updateTaskStatus(task.id, 'rendering');
    }
    
    try {
      console.log(`[AIQueue] 开始处理任务: ${task.id}`);
      
      // 调用Gemini API
      const result = await callGeminiAPIWithPromise(
        task.image1,
        task.image2,
        task.image1Width,
        task.image1Height,
        task.additionalPrompt || ''
      );
      
      if (result.success && result.imageData) {
        // 保存图片到用户目录
        const imageUrl = await saveAIImage(task.username, task.id, result.imageData);
        
        // 更新任务状态
        task.status = 'completed';
        task.imageUrl = imageUrl;
        task.completedAt = new Date().toISOString();
        updateTaskStatus(task.id, 'completed', imageUrl);
        
        console.log(`[AIQueue] 任务完成: ${task.id}`);
      } else {
        // Gemini API 明确返回失败
        const apiError = new Error(result.error || '生成失败');
        apiError.isApiError = true;
        throw apiError;
      }
    } catch (error) {
      console.error(`[AIQueue] 任务失败: ${task.id}`, error);
      
      // 只有 Gemini API 明确返回失败时才标记为 failed
      if (error.isApiError) {
        task.status = 'failed';
        task.error = error.message;
        task.completedAt = new Date().toISOString();
        // 保存原始任务数据用于重试（不包含图片数据以节省空间，重试时从预览图重新加载）
        updateTaskStatus(task.id, 'failed', null, error.message, {
          image1Width: task.image1Width,
          image1Height: task.image1Height,
          additionalPrompt: task.additionalPrompt
        });
      } else {
        // 其他错误（代码错误、网络错误等）自动重试
        console.log(`[AIQueue] 非API错误，将任务重新加入队列末尾: ${task.id}`);
        task.status = 'queued';
        task.retryCount = (task.retryCount || 0) + 1;
        
        // 最多重试3次
        if (task.retryCount <= 3) {
          queue.push({ ...task });
          updateTaskStatus(task.id, 'queued');
        } else {
          console.error(`[AIQueue] 任务重试次数超过限制，标记为失败: ${task.id}`);
          task.status = 'failed';
          task.error = '多次重试失败：' + error.message;
          task.completedAt = new Date().toISOString();
          updateTaskStatus(task.id, 'failed', null, task.error, {
            image1Width: task.image1Width,
            image1Height: task.image1Height,
            additionalPrompt: task.additionalPrompt
          });
        }
      }
    }
    
    // 从队列中移除
    queue.shift();
    saveQueue();
  }
  
  isProcessing = false;
}

// 调用Gemini API（Promise版本）
function callGeminiAPIWithPromise(image1Base64, image2Base64, image1Width, image1Height, additionalPrompt) {
  return new Promise((resolve, reject) => {
    const https = require('https');
    
    // 移除 data:image/png;base64, 前缀（如果有）
    const cleanImage1 = image1Base64.replace(/^data:image\/\w+;base64,/, '');
    const cleanImage2 = image2Base64.replace(/^data:image\/\w+;base64,/, '');
    
    // 构建请求内容
    const content = [
      {
        type: "image_url",
        image_url: {
          url: `data:image/png;base64,${cleanImage1}`
        }
      },
      {
        type: "image_url",
        image_url: {
          url: `data:image/png;base64,${cleanImage2}`
        }
      },
      {
        type: "text",
        text: ReplaceCharacterHandler.buildPromptText(image1Width, image1Height, additionalPrompt)
      }
    ];
    
    const requestData = JSON.stringify({
      model: "gemini-3-pro-image-preview",
      messages: [{ role: "user", content: content }]
    });
    
    const options = {
      hostname: 'api.chatanywhere.tech',
      port: 443,
      path: '/v1/chat/completions',
      method: 'POST',
      headers: {
        'Authorization': 'Bearer sk-j32LgDixK6pfESYGfJtgc2Tzlmszx5NZhSH0sOzpLQkYuKek',
        'Content-Type': 'application/json',
        'Content-Length': Buffer.byteLength(requestData)
      },
      timeout: 300000 // 5分钟超时
    };
    
    console.log('[AIQueue] 正在调用 Gemini API...');
    
    const geminiReq = https.request(options, (geminiRes) => {
      let responseData = '';
      
      geminiRes.on('data', (chunk) => {
        responseData += chunk;
      });
      
      geminiRes.on('end', () => {
        try {
          if (geminiRes.statusCode !== 200) {
            console.error('[AIQueue] Gemini API 返回错误:', geminiRes.statusCode, responseData);
            reject(new Error(`Gemini API error: ${geminiRes.statusCode}`));
            return;
          }
          
          const response = JSON.parse(responseData);
          
          // 解析响应，提取图片
          const imageData = ReplaceCharacterHandler.extractImageFromResponse(response);
          
          if (!imageData) {
            console.error('[AIQueue] 无法从响应中提取图片');
            reject(new Error('Failed to extract image from response'));
            return;
          }
          
          resolve({
            success: true,
            imageData: imageData
          });
          
          console.log('[AIQueue] ✓ 成功处理请求');
        } catch (error) {
          console.error('[AIQueue] 解析响应失败:', error);
          reject(error);
        }
      });
    });
    
    geminiReq.on('error', (error) => {
      console.error('[AIQueue] 请求错误:', error);
      reject(error);
    });
    
    geminiReq.on('timeout', () => {
      console.error('[AIQueue] 请求超时');
      geminiReq.destroy();
      reject(new Error('Request timeout'));
    });
    
    geminiReq.write(requestData);
    geminiReq.end();
  });
}

// 保存AI生成的图片
function saveAIImage(username, taskId, imageBase64) {
  const usersDir = path.join(__dirname, 'users');
  const userDir = path.join(usersDir, username.toLowerCase());
  const aiDir = path.join(userDir, 'ai-images');
  
  // 确保目录存在
  if (!fs.existsSync(aiDir)) {
    fs.mkdirSync(aiDir, { recursive: true });
  }
  
  const imagePath = path.join(aiDir, `${taskId}.png`);
  const imageBuffer = Buffer.from(imageBase64, 'base64');
  fs.writeFileSync(imagePath, imageBuffer);
  
  return `/api/ai/image?username=${encodeURIComponent(username)}&id=${encodeURIComponent(taskId)}`;
}

// 更新任务状态
function updateTaskStatus(taskId, status, imageUrl = null, error = null, retryData = null) {
  const history = loadAIHistory();
  
  for (const username in history) {
    const userHistory = history[username];
    const task = userHistory.find(t => t.id === taskId);
    if (task) {
      task.status = status;
      if (imageUrl) {
        task.imageUrl = imageUrl;
      }
      if (error) {
        task.error = error;
      }
      if (retryData) {
        task.retryData = retryData; // 保存重试所需数据
      }
      if (status === 'completed' || status === 'failed') {
        task.completedAt = new Date().toISOString();
      }
      break;
    }
  }
  
  saveAIHistory(history);
}

// 获取用户AI历史
function getUserAIHistory(username) {
  const history = loadAIHistory();
  const normalizedUsername = username.toLowerCase();
  return history[normalizedUsername] || [];
}

// 处理AI生图请求（队列版本）
function handleAIRequest(req, res) {
  if (req.method !== 'POST') {
    res.writeHead(405, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ error: 'Method not allowed' }));
    return;
  }

  let body = '';
  
  req.on('data', (chunk) => {
    body += chunk.toString();
  });

  req.on('end', () => {
    try {
      const data = JSON.parse(body);
      const { username, image1, image2, image1Width, image1Height, additionalPrompt } = data;

      if (!username) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ success: false, error: '缺少用户名参数' }));
        return;
      }

      if (!image1 || !image2) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ success: false, error: 'Missing required fields: image1, image2' }));
        return;
      }

      if (!image1Width || !image1Height) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ success: false, error: 'Missing required fields: image1Width, image1Height' }));
        return;
      }

      // 添加到队列
      const taskId = addToQueue(username, {
        image1,
        image2,
        image1Width,
        image1Height,
        additionalPrompt: additionalPrompt || ''
      });

      // 立即返回任务ID
      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({
        success: true,
        taskId: taskId,
        message: '请求生图成功，正在处理中...'
      }));

      // 异步处理队列
      processQueue();
    } catch (error) {
      console.error('[AIQueue] 处理请求失败:', error);
      res.writeHead(500, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ success: false, error: '处理失败', details: error.message }));
    }
  });

  req.on('error', (error) => {
    console.error('[AIQueue] 请求错误:', error);
    res.writeHead(500, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ success: false, error: 'Request error', details: error.message }));
  });
}

// 重试失败的任务（免费）
function retryTask(taskId, username) {
  const history = loadAIHistory();
  const userHistory = history[username.toLowerCase()];
  
  if (!userHistory) {
    return { success: false, error: '用户历史不存在' };
  }
  
  const task = userHistory.find(t => t.id === taskId);
  if (!task) {
    return { success: false, error: '任务不存在' };
  }
  
  if (task.status !== 'failed') {
    return { success: false, error: '只能重试失败的任务' };
  }
  
  // 读取预览图作为 image1
  const previewPath = path.join(__dirname, 'users', username.toLowerCase(), 'ai-images', `${taskId}_preview.png`);
  if (!fs.existsSync(previewPath)) {
    return { success: false, error: '预览图不存在，无法重试' };
  }
  
  const previewData = fs.readFileSync(previewPath);
  const image1Base64 = previewData.toString('base64');
  
  // 生成新任务ID
  const newTaskId = `ai_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  
  // 获取重试数据
  const retryData = task.retryData || {};
  
  // 创建新任务
  const newTask = {
    id: newTaskId,
    username: username.toLowerCase(),
    status: queue.length === 0 && !isProcessing ? 'rendering' : 'queued',
    createdAt: new Date().toISOString(),
    image1: image1Base64,
    image2: image1Base64, // 重试时使用同一张图
    image1Width: retryData.image1Width || 512,
    image1Height: retryData.image1Height || 512,
    additionalPrompt: retryData.additionalPrompt || '',
    isRetry: true,
    originalTaskId: taskId
  };
  
  queue.push(newTask);
  saveQueue();
  
  // 保存预览图到新任务
  const newPreviewPath = path.join(__dirname, 'users', username.toLowerCase(), 'ai-images', `${newTaskId}_preview.png`);
  fs.copyFileSync(previewPath, newPreviewPath);
  
  // 添加到历史记录
  userHistory.unshift({
    id: newTaskId,
    status: newTask.status,
    createdAt: newTask.createdAt,
    imageUrl: null,
    previewUrl: `/api/ai/preview?username=${encodeURIComponent(username)}&id=${encodeURIComponent(newTaskId)}`,
    isRetry: true
  });
  
  // 将原任务标记为已重试
  task.retried = true;
  task.retriedTaskId = newTaskId;
  
  saveAIHistory(history);
  
  // 开始处理队列
  if (queue.length === 1 && !isProcessing) {
    processQueue();
  }
  
  return { success: true, taskId: newTaskId };
}

// 处理重试请求
function handleRetryRequest(req, res) {
  let body = '';
  
  req.on('data', chunk => {
    body += chunk.toString();
  });
  
  req.on('end', () => {
    try {
      const { taskId, username } = JSON.parse(body);
      
      if (!taskId || !username) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ success: false, error: '缺少必要参数' }));
        return;
      }
      
      const result = retryTask(taskId, username);
      
      res.writeHead(result.success ? 200 : 400, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify(result));
    } catch (error) {
      console.error('[AIQueue] 重试请求失败:', error);
      res.writeHead(500, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ success: false, error: '处理失败', details: error.message }));
    }
  });
}

// 初始化
initQueue();

module.exports = {
  handleAIRequest,
  handleRetryRequest,
  getUserAIHistory,
  processQueue
};