gongxiangp2/yedir/server/utils/image-utils.js

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<Object>} 保存的图片信息
 */
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<Object>} 保存结果
 */
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<Object|null>} 图片信息或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<Array>} 图片信息列表
 */
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<Array>} 图片信息列表
 */
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
};