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');

// 超时配置（毫秒）
const TASK_TIMEOUT = 5 * 60 * 1000; // 5分钟超时

// 队列状态
let queue = [];
let isProcessing = false;
let dbInstance = null;
let timeoutCheckInterval = 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 checkTimeoutTasks() {
  try {
    const db = await getDB();
    const timedOutTasks = db.checkAndMarkTimeoutTasks(TASK_TIMEOUT);

    if (timedOutTasks && timedOutTasks.length > 0) {
      console.log(`[AIQueue] 发现 ${timedOutTasks.length} 个超时任务，已标记为失败`);
      timedOutTasks.forEach(taskId => {
        console.log(`[AIQueue] 超时任务: ${taskId}`);
      });
    }
  } catch (error) {
    console.error('[AIQueue] 检查超时任务失败:', error);
  }
}

// 启动超时检查定时器
function startTimeoutChecker() {
  if (timeoutCheckInterval) {
    clearInterval(timeoutCheckInterval);
  }
  // 每30秒检查一次超时任务
  timeoutCheckInterval = setInterval(checkTimeoutTasks, 30000);
  console.log('[AIQueue] 超时检查定时器已启动（每30秒检查一次）');
}

// 处理队列
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';
      task.renderStartTime = Date.now();
      updateTaskStatus(task.id, 'rendering', null, null, null, task.renderStartTime);
    }
    
    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');
    const startTime = Date.now();

    // 移除 data:image/png;base64, 前缀（如果有）
    const cleanImage1 = image1Base64.replace(/^data:image\/\w+;base64,/, '');
    const cleanImage2 = image2Base64.replace(/^data:image\/\w+;base64,/, '');

    // 构建请求内容
    const promptText = ReplaceCharacterHandler.buildPromptText(image1Width, image1Height, additionalPrompt);
    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: promptText
      }
    ];

    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('\n========== [Gemini API 请求] ==========');
    console.log(`[时间] ${new Date().toLocaleString()}`);
    console.log(`[模型] gemini-3-pro-image-preview`);
    console.log(`[目标] https://${options.hostname}${options.path}`);
    console.log(`[图片1大小] ${(cleanImage1.length / 1024).toFixed(2)} KB`);
    console.log(`[图片2大小] ${(cleanImage2.length / 1024).toFixed(2)} KB`);
    console.log(`[输出尺寸] ${image1Width} x ${image1Height}`);
    console.log(`[附加提示] ${additionalPrompt || '(无)'}`);
    console.log(`[请求体大小] ${(Buffer.byteLength(requestData) / 1024).toFixed(2)} KB`);
    console.log('========================================\n');

    const geminiReq = https.request(options, (geminiRes) => {
      let responseData = '';

      geminiRes.on('data', (chunk) => {
        responseData += chunk;
      });

      geminiRes.on('end', () => {
        const elapsed = ((Date.now() - startTime) / 1000).toFixed(2);

        console.log('\n========== [Gemini API 响应] ==========');
        console.log(`[时间] ${new Date().toLocaleString()}`);
        console.log(`[耗时] ${elapsed} 秒`);
        console.log(`[状态码] ${geminiRes.statusCode}`);
        console.log(`[响应大小] ${(responseData.length / 1024).toFixed(2)} KB`);

        try {
          if (geminiRes.statusCode !== 200) {
            console.log(`[错误] API返回非200状态`);
            console.log(`[响应内容] ${responseData.substring(0, 500)}...`);
            console.log('========================================\n');
            reject(new Error(`Gemini API error: ${geminiRes.statusCode}`));
            return;
          }

          const response = JSON.parse(responseData);

          // 打印响应结构
          console.log(`[响应结构] choices: ${response.choices?.length || 0}`);
          if (response.choices && response.choices[0]) {
            const msg = response.choices[0].message;
            if (msg && msg.content) {
              if (Array.isArray(msg.content)) {
                console.log(`[内容类型] 数组，包含 ${msg.content.length} 个元素`);
                msg.content.forEach((item, i) => {
                  if (item.type === 'text') {
                    console.log(`  [${i}] text: ${item.text?.substring(0, 100) || '(空)'}...`);
                  } else if (item.type === 'image_url') {
                    const imgData = item.image_url?.url || '';
                    console.log(`  [${i}] image: ${(imgData.length / 1024).toFixed(2)} KB`);
                  } else {
                    console.log(`  [${i}] ${item.type}: ...`);
                  }
                });
              } else {
                console.log(`[内容类型] 字符串，长度 ${msg.content.length}`);
              }
            }
          }
          if (response.usage) {
            console.log(`[Token使用] prompt: ${response.usage.prompt_tokens}, completion: ${response.usage.completion_tokens}, total: ${response.usage.total_tokens}`);
          }

          // 解析响应，提取图片
          const imageData = ReplaceCharacterHandler.extractImageFromResponse(response);

          if (!imageData) {
            console.log(`[结果] ✗ 无法从响应中提取图片`);
            console.log('========================================\n');
            reject(new Error('Failed to extract image from response'));
            return;
          }

          console.log(`[结果] ✓ 成功提取图片，大小: ${(imageData.length / 1024).toFixed(2)} KB`);
          console.log('========================================\n');

          resolve({
            success: true,
            imageData: imageData
          });
        } catch (error) {
          console.log(`[错误] 解析响应失败: ${error.message}`);
          console.log('========================================\n');
          reject(error);
        }
      });
    });

    geminiReq.on('error', (error) => {
      const elapsed = ((Date.now() - startTime) / 1000).toFixed(2);
      console.log('\n========== [Gemini API 错误] ==========');
      console.log(`[时间] ${new Date().toLocaleString()}`);
      console.log(`[耗时] ${elapsed} 秒`);
      console.log(`[错误类型] ${error.code || 'Unknown'}`);
      console.log(`[错误信息] ${error.message}`);
      console.log('========================================\n');
      reject(error);
    });

    geminiReq.on('timeout', () => {
      const elapsed = ((Date.now() - startTime) / 1000).toFixed(2);
      console.log('\n========== [Gemini API 超时] ==========');
      console.log(`[时间] ${new Date().toLocaleString()}`);
      console.log(`[耗时] ${elapsed} 秒`);
      console.log(`[超时设置] ${options.timeout / 1000} 秒`);
      console.log('========================================\n');
      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, renderStartTime = null) {
  try {
    const db = await getDB();
    db.updateAITaskStatus(taskId, status, imageUrl, error, renderStartTime);
    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();
startTimeoutChecker();

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