const mongoose = require('mongoose'); const GeneratedImage = require('../models/GeneratedImage'); const path = require('path'); const fs = require('fs'); const { v4: uuidv4 } = require('uuid'); const sharp = require('sharp'); const fetch = require('node-fetch'); const logger = require('./logger'); const User = require('../models/User'); const { uploadToMinio } = require('./minio-client'); /** * 保存生成的图片 * @param {Object} options - 保存选项 * @param {Buffer|string} options.imageData - 图片数据Buffer或图片URL * @param {string} options.userId - 用户ID * @param {string} options.prompt - 生成图片的提示词 * @param {string} options.type - 图片类型 (如 'dalle', 'sd', 'midjourney') * @param {Object} options.metadata - 额外的元数据 * @param {boolean} options.saveToDb - 是否保存到数据库,默认为true * @returns {Promise} 保存的图片信息 */ async function saveGeneratedImage(options) { const { imageData, userId, prompt = '', type = 'unknown', metadata = {}, saveToDb = true } = options; if (!imageData) { throw new Error('未提供图片数据'); } // 生成唯一文件名和路径 const timestamp = Date.now(); const filename = `${timestamp}-${uuidv4()}.png`; const relativePath = `downloads/${type}/${filename}`; const fullPath = path.join(__dirname, '../../public/uploads', relativePath); const webPath = `/uploads/${relativePath}`; // 确保目录存在 const dir = path.dirname(fullPath); if (!fs.existsSync(dir)) { fs.mkdirSync(dir, { recursive: true }); } let buffer; // 获取图片数据 if (Buffer.isBuffer(imageData)) { buffer = imageData; } else if (typeof imageData === 'string' && (imageData.startsWith('http://') || imageData.startsWith('https://'))) { // 如果是URL,下载图片 try { const response = await fetch(imageData); if (!response.ok) { throw new Error(`获取图片失败: ${response.statusText}`); } buffer = await response.buffer(); } catch (error) { throw new Error(`下载图片失败: ${error.message}`); } } else if (typeof imageData === 'string' && imageData.startsWith('data:image/')) { // 如果是Base64格式的数据URL const base64Data = imageData.replace(/^data:image\/\w+;base64,/, ''); buffer = Buffer.from(base64Data, 'base64'); } else { throw new Error('不支持的图片数据格式'); } // 保存图片到本地文件系统 try { await sharp(buffer).png().toFile(fullPath); } catch (error) { throw new Error(`保存图片失败: ${error.message}`); } // 如果启用了S3存储,上传到S3 let minioUrl = null; if (process.env.USE_S3_STORAGE === 'true') { try { // 修改:确保S3存储路径包含downloads前缀 const s3Key = `downloads/${type}/${filename}`; minioUrl = await uploadToMinio({ buffer: buffer, mimetype: 'image/png', originalname: filename, s3Key: s3Key // 显式设置S3键,包含downloads路径 }); logger.info(`图片已上传到S3,URL: ${minioUrl}`); } catch (error) { logger.error(`上传图片到S3失败: ${error.message}`); // 如果S3上传失败,继续使用本地存储路径 } } // 返回结果 const result = { path: fullPath, webPath: minioUrl || webPath, url: minioUrl || webPath }; // 保存到数据库 if (saveToDb && userId) { try { // 检查用户是否存在 const user = await User.findById(userId); if (!user) { logger.warn(`用户不存在,ID: ${userId},跳过数据库保存`); } else { const generatedImage = new GeneratedImage({ user: userId, prompt, imageUrl: minioUrl || webPath, originalImageUrl: minioUrl || webPath, type, metadata }); await generatedImage.save(); result.id = generatedImage._id; logger.info(`已将图片保存到数据库,ID: ${generatedImage._id}`); } } catch (error) { logger.error(`保存图片到数据库失败: ${error.message}`); // 继续执行,即使数据库保存失败 } } return result; } /** * 保存API返回的图片结果到用户的图片库 * @param {string} imageUrl - API返回的图片URL * @param {string} userId - 用户ID * @param {string} type - 图片类型 * @param {string} prompt - 可选的提示词 * @param {Object} metadata - 可选的元数据 * @returns {Promise} 保存结果 */ async function saveApiResultImage(imageUrl, userId, type = 'other', prompt = '', metadata = {}) { try { // 确保用户ID存在 if (!userId) { throw new Error('未提供用户ID'); } // 确保图片URL存在 if (!imageUrl) { throw new Error('未提供图片URL'); } logger.info(`准备保存API结果图片: ${imageUrl} 到用户 ${userId} 的图片库`); // 下载图片 let response; try { response = await fetch(imageUrl); if (!response.ok) { throw new Error(`获取图片失败: ${response.status} ${response.statusText}`); } } catch (error) { logger.error(`下载图片失败: ${error.message}`); throw error; } // 获取图片buffer const imageBuffer = await response.buffer(); // 使用saveGeneratedImage保存图片 const savedImage = await saveGeneratedImage({ imageData: imageBuffer, userId, prompt, type, metadata, saveToDb: true }); logger.info(`API结果图片成功保存: ${savedImage.url}`); return savedImage; } catch (error) { logger.error(`保存API结果图片失败: ${error.message}`); throw error; } } /** * 根据任务ID获取生成的图片信息 * @param {string} taskId - 任务ID * @returns {Promise} 图片信息或null */ async function getGeneratedImageByTaskId(taskId) { try { return await GeneratedImage.findOne({ taskId }); } catch (error) { console.error('根据任务ID获取生成的图片信息失败:', error); return null; } } /** * 根据用户ID和类型获取生成的图片列表 * @param {string} userId - 用户ID * @param {string} type - 图片类型 * @param {number} [limit=10] - 限制返回的数量 * @returns {Promise} 图片信息列表 */ async function getGeneratedImagesByUserAndType(userId, type, limit = 10) { try { return await GeneratedImage.find({ user: userId, type }) .sort({ createdAt: -1 }) .limit(limit); } catch (error) { console.error('获取用户生成的图片列表失败:', error); return []; } } /** * 根据用户ID获取所有生成的图片列表 * @param {string} userId - 用户ID * @param {number} [limit=20] - 限制返回的数量 * @returns {Promise} 图片信息列表 */ async function getAllGeneratedImagesByUser(userId, limit = 20) { try { return await GeneratedImage.find({ user: userId }) .sort({ createdAt: -1 }) .limit(limit); } catch (error) { console.error('获取用户所有生成的图片列表失败:', error); return []; } } module.exports = { saveGeneratedImage, saveApiResultImage, getGeneratedImageByTaskId, getGeneratedImagesByUserAndType, getAllGeneratedImagesByUser };