AnimationManager/server/ai-queue.js

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

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

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

// 获取数据库实例
async function getDB() {
  if (!dbInstance) {
    dbInstance = await getDatabase();
  }
  return dbInstance;
}

// 初始化：加载队列
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);
  }
}

// 添加任务到队列
async 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();
  
  // 添加到历史记录（不保存参考图）
  try {
    const db = await getDB();
    db.addAIHistory(taskId, username, task.status, null);
  } catch (error) {
    console.error('[AIQueue] 添加历史记录失败:', error);
  }
  
  // 如果队列为空且没有正在处理的任务，立即开始处理
  if (queue.length === 1 && !isProcessing) {
    processQueue();
  }
  
  return 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)}`;
}

// 更新任务状态
async function updateTaskStatus(taskId, status, imageUrl = null, error = null, retryData = null) {
  try {
    const db = await getDB();
    db.updateAITaskStatus(taskId, status, imageUrl, error);
    if (retryData) {
      db.updateAITaskRetryData(taskId, retryData);
    }
  } catch (err) {
    console.error('[AIQueue] 更新任务状态失败:', err);
  }
}

// 获取用户AI历史
async function getUserAIHistory(username) {
  try {
    const db = await getDB();
    return db.getAIHistory(username);
  } catch (error) {
    console.error('[AIQueue] 获取用户AI历史失败:', error);
    return [];
  }
}

// 处理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', async () => {
    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 = await 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 }));
  });
}

// 重试失败的任务 - 直接删除失败记录（不再支持重试，因为不保存参考图）
async function retryTask(taskId, username) {
  try {
    const db = await getDB();
    const task = db.getAITask(taskId);
    
    if (!task) {
      return { success: false, error: '任务不存在' };
    }
    
    if (task.status !== 'failed') {
      return { success: false, error: '只能删除失败的任务' };
    }
    
    // 直接删除失败的任务记录
    db.deleteAITask(taskId);
    
    return { success: true, message: '已删除失败记录' };
  } catch (error) {
    console.error('[AIQueue] 删除任务失败:', error);
    return { success: false, error: error.message };
  }
}

// 处理重试请求
function handleRetryRequest(req, res) {
  let body = '';
  
  req.on('data', chunk => {
    body += chunk.toString();
  });
  
  req.on('end', async () => {
    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 = await 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
};