ai-swap/backend/services/ai_swap_cloth_service.py

import os
import uuid
import random
import hashlib
import numpy as np
from PIL import Image
from datetime import datetime
from typing import Dict, Any, Optional, Tuple

from backend.modules.comfyui.ai_swap_cloth import ai_swap_cloth_process
from backend.modules.comfyui.ai_copywriter import gen_copywriter
from backend.modules.database.operations import DatabaseOperations
from backend.modules.database.models import ProcessRecord, ImageRecord, User
from backend.modules.database.connection import DatabaseConnection
from backend.modules.database.models import SystemConfig
from backend.utils.logger_config import setup_logger
from backend.utils.system_config import Config
from backend.services.task_queue_service import get_task_queue_service, TaskStatus

logger = setup_logger(__name__)

class AISwapClothService:
    """
    AI换衣服业务逻辑服务
    负责协调AI处理、文案生成、数据库操作等完整业务流程
    """

    def __init__(self, db_operations: Optional[DatabaseOperations] = None):
        """
        初始化服务

        Args:
            db_operations: 数据库操作对象，如果为None则创建默认实例
        """
        self.db_ops = db_operations or DatabaseOperations()
        self.system_config = Config('./backend/config/ai_swap_cloth_config.json')
        self.task_queue = get_task_queue_service()

        # 确保输出目录存在
        os.makedirs(self.system_config.output_dir, exist_ok=True)

    def submit_swap_cloth_task(
        self, 
        user_id: int,
        raw_image_id: int,
        cloth_image_id: int,
        quantity: int = 1
    ) -> str:
        """
        提交换衣服任务到队列（非阻塞）

        Args:
            user_id: 用户ID
            raw_image_id: 原始图片ID
            cloth_image_id: 衣服图片ID
            quantity: 生成数量（默认1）

        Returns:
            str: 任务ID
        """
        logger.info(f"提交用户 {user_id} 的换衣服任务到队列")

        # 提交任务到队列
        task_id = self.task_queue.submit_task(
            task_func=self._execute_swap_cloth_task,
            task_args=(user_id, raw_image_id, cloth_image_id),
            task_kwargs={
                "quantity": max(1, int(quantity))
            }
        )

        logger.info(f"任务已提交到队列，任务ID: {task_id}")
        return task_id

    def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """
        获取任务状态
        
        Args:
            task_id: 任务ID
            
        Returns:
            Dict: 任务状态信息
        """
        return self.task_queue.get_task_status(task_id)
    
    def get_user_tasks(self, user_id: int) -> list:
        """
        获取用户的所有任务
        
        Args:
            user_id: 用户ID
            
        Returns:
            list: 任务列表
        """
        return self.task_queue.get_all_tasks(user_id)
    
    def cancel_task(self, task_id: str) -> bool:
        """
        取消任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功取消
        """
        return self.task_queue.cancel_task(task_id)

    def _execute_swap_cloth_task(
        self,
        user_id: int,
        raw_image_id: int,
        cloth_image_id: int,
        quantity: int = 1
    ) -> Dict[str, Any]:
        """
        执行AI换衣服任务（在线程池中运行）

        Args:
            user_id: 用户ID
            raw_image_id: 原始图片ID
            cloth_image_id: 衣服图片ID
            quantity: 生成数量（默认1）

        Returns:
            Dict: 包含处理结果的字典
        """
        try:
            logger.info(f"开始执行用户 {user_id} 的换衣服任务")

            # 1. 输入验证
            self._validate_inputs(user_id, raw_image_id, cloth_image_id)

            # 2. 获取输入图片
            raw_image, cloth_image = self._get_input_images(raw_image_id, cloth_image_id)

            # 3. 循环生成 quantity 次
            total_count = max(1, int(quantity))
            process_record_ids = []
            result_image_ids = []
            copywriter_texts = []
            history_prompt_last = None

            # 为了提升多张图的差异性，循环内随机化种子（若底层支持）
            try:
                from backend.modules.comfyui import ai_swap_cloth as comfy_ai_swap_cloth_module
            except Exception:
                comfy_ai_swap_cloth_module = None

            for _ in range(total_count):
                if comfy_ai_swap_cloth_module is not None:
                    try:
                        comfy_ai_swap_cloth_module.system_config.seed = random.randint(1, 10_000_000)
                    except Exception:
                        pass

                # 执行AI换衣服
                result_image, history_prompt = self._process_ai_swap_cloth(
                    raw_image, cloth_image
                )
                history_prompt_last = history_prompt

                # 生成文案描述
                copywriter_text = self._generate_copywriter(result_image)
                copywriter_texts.append(copywriter_text)

                # 保存结果图片
                result_image_record = self._save_result_image(user_id, result_image)
                result_image_ids.append(result_image_record["id"])

                # 创建处理记录
                process_record = self._create_process_record(
                    user_id, raw_image_id, cloth_image_id, 
                    result_image_record["id"], copywriter_text, prompt="this is the prompt of swap cloth task"
                )
                process_record_ids.append(process_record["id"])

            logger.info(
                f"用户 {user_id} 的换衣服任务完成，共生成 {len(process_record_ids)} 张，首个记录ID: {process_record_ids[0]}"
            )

            # 为兼容旧前端，保留单值字段，同时返回批量字段
            return {
                "success": True,
                "count": len(process_record_ids),
                "process_record_id": process_record_ids[0],
                "result_image_id": result_image_ids[0],
                "copywriter_text": copywriter_texts[0] if copywriter_texts else None,
                "history_prompt": history_prompt_last,
                "process_record": self.db_ops.get_process_record_by_id(process_record_ids[0]),
                "process_record_ids": process_record_ids,
                "result_image_ids": result_image_ids,
                "copywriter_texts": copywriter_texts,
            }
            
        except Exception as e:
            logger.error(f"换衣服任务执行失败: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__
            }
    
    def process_swap_cloth_with_record(
        self,
        user_id: int,
        raw_image_id: int,
        cloth_image_id: int,
        **kwargs
    ) -> Dict[str, Any]:
        """
        完整的AI换衣服业务流程（同步版本，保持向后兼容）
        
        Args:
            user_id: 用户ID
            raw_image_id: 原始图片ID
            cloth_image_id: 衣服图片ID
            **kwargs: 其他可选参数

        Returns:
            Dict: 包含处理结果的字典
        """
        return self._execute_swap_cloth_task(
            user_id, raw_image_id, cloth_image_id, **kwargs
        )
    
    def _validate_inputs(self, user_id: int, raw_image_id: int, cloth_image_id: int):
        """
        验证输入参数

        Args:
            user_id: 用户ID
            raw_image_id: 原始图片ID
            cloth_image_id: 衣服图片ID
        """
        # 验证用户是否存在
        user = self.db_ops.get_user_by_id(user_id)
        if not user:
            raise ValueError(f"用户ID {user_id} 不存在")
        
        if not user.get("is_active", False):
            raise ValueError(f"用户ID {user_id} 已被禁用")

        # 验证原始图片
        raw_image = self.db_ops.get_image_record_by_id(raw_image_id)
        if not raw_image:
            raise ValueError(f"原始图片ID {raw_image_id} 不存在")
        
        if raw_image["image_type"] != "original":
            raise ValueError(f"原始图片ID {raw_image_id} 不是原始图片类型")
        
        if raw_image["user_id"] != user_id:
            raise ValueError(f"原始图片ID {raw_image_id} 不属于用户 {user_id}")

        # 验证衣服图片
        cloth_image = self.db_ops.get_image_record_by_id(cloth_image_id)
        if not cloth_image:
            raise ValueError(f"衣服图片ID {cloth_image_id} 不存在")
        
        if cloth_image["image_type"] != "cloth":
            raise ValueError(f"衣服图片ID {cloth_image_id} 不是衣服图片类型")

        if cloth_image["user_id"] != user_id:
            raise ValueError(f"衣服图片ID {cloth_image_id} 不属于用户 {user_id}")

        logger.info(f"输入验证通过: 用户={user_id}, 原始图片={raw_image_id}, 衣服图片={cloth_image_id}")
    
    def _get_input_images(self, raw_image_id: int, cloth_image_id: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        获取输入图片数据
        
        Args:
            raw_image_id: 原始图片ID
            cloth_image_id: 衣服图片ID

        Returns:
            Tuple: (原始图片数组, 衣服图片数组)
        """
        # 获取原始图片
        raw_image = self.db_ops.get_image_record_by_id(raw_image_id)
        if not os.path.exists(raw_image["stored_path"]):
            raise FileNotFoundError(f"原始图片文件不存在: {raw_image['stored_path']}")
        
        raw_image = np.array(Image.open(raw_image["stored_path"]))

        # 获取衣服图片
        cloth_record = self.db_ops.get_image_record_by_id(cloth_image_id)
        if not os.path.exists(cloth_record["stored_path"]):
            raise FileNotFoundError(f"衣服图片文件不存在: {cloth_record['stored_path']}")

        cloth_image = np.array(Image.open(cloth_record["stored_path"]))

        logger.info(f"成功加载输入图片: 原始图片={raw_image.shape}, 衣服图片={cloth_image.shape}")

        return raw_image, cloth_image
    
    def _process_ai_swap_cloth(
        self,
        raw_image: np.ndarray,
        cloth_image: np.ndarray
    ) -> Tuple[Image.Image, str]:
        """
        执行AI换衣服处理

        Args:
            raw_image: 原始图片数组
            cloth_image: 衣服图片数组

        Returns:
            Tuple: (结果图片， 历史提示词)
        """
        logger.info(f"开始执行AI换衣服处理，原始图片={raw_image.shape}, 衣服图片={cloth_image.shape}")

        try:
            result_image, history_prompt = ai_swap_cloth_process(
                raw_image, cloth_image
            )

            logger.info("AI换衣服处理完成")
            return result_image, history_prompt

        except Exception as e:
            logger.error(f"AI换衣服处理失败: {str(e)}", exc_info=True)

    def _generate_copywriter(self, result_image: Image.Image) -> str:
        """
        基于结果图片生成文案描述
        
        Args:
            result_image: 结果图片
            
        Returns: 
            str: 生成的文案描述
        """
        logger.info("开始生成文案描述")

        try:
            copywriter_text = gen_copywriter(result_image)
            logger.info("文案描述生成完成")
            return copywriter_text

        except Exception as e:
            logger.error(f"文案生成失败: {str(e)}", exc_info=True)
            # 文案生成失败不影响主流程，返回默认文案
            return "AI换衣服完成，效果很棒！✨"

    def _save_result_image(self, user_id: int, result_image: Image.Image) -> Dict[str, Any]:
        """
        保存结果图片到数据库和文件系统
        
        Args:
            user_id: 用户ID
            result_image: 结果图片
            
        Returns:
            Dict: 图片记录信息
        """
        # 生成唯一文件名
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"result_{user_id}_{timestamp}_{unique_id}.png"
        
        # 保存到文件系统
        file_path = os.path.join(self.system_config.output_dir, filename)
        result_image.save(file_path, "PNG")
        
        # 计算文件大小和哈希值
        file_size = os.path.getsize(file_path)
        image_hash = self._calculate_image_hash(file_path)
        
        # 保存到数据库
        image_record = self.db_ops.create_image_record(
            user_id=user_id,
            image_type="result",
            original_filename=filename,
            stored_path=file_path,
            file_size=file_size,
            image_hash=image_hash
        )
        
        logger.info(f"结果图片保存成功: {file_path}, 记录ID: {image_record['id']}")
        return image_record

    def _create_process_record(
        self,
        user_id: int,
        raw_image_id: int,
        cloth_image_id: int,
        result_image_id: int,
        copywriter_text: str,
        prompt: str
    ) -> Dict[str, Any]:
        """
        创建处理记录
        
        Args:
            user_id: 用户ID
            raw_image_id: 原始图片ID
            cloth_image_id: 衣服图片ID
            result_image_id: 结果图片ID
            copywriter_text: 文案描述
            prompt: 提示词
        Returns:
            Dict: 处理记录信息
        """
        process_record = self.db_ops.create_process_record(
            user_id=user_id,
            face_image_id=raw_image_id,
            cloth_image_id=cloth_image_id,
            result_image_id=result_image_id,
            generated_text=copywriter_text,
            task_type="swap_cloth",
            prompt=prompt
        )

        # 更新完成时间
        self.db_ops.update_process_record(
            process_record["id"],
            {"completed_at": datetime.now()}
        )

        logger.info(f"处理记录创建成功: {process_record['id']}")
        return process_record
    
    def _calculate_image_hash(self, image_path: str) -> str:
        """
        计算图片哈希值
        
        Args:
            image_path: 图片路径
            
        Returns:
            str: MD5哈希值
        """
        with open(image_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()

    def get_user_process_history(self, user_id: int, page: int = 1, page_size: int = 20) -> Dict[str, Any]:
        """
        获取用户的处理历史记录
        
        Args:
            user_id: 用户ID
            page: 页码
            page_size: 每页大小
            
        Returns:
            Dict: 分页的处理记录列表
        """
        return self.db_ops.get_user_process_records(user_id, page, page_size)
    
    def get_process_detail(self, process_id: int, user_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
        """
        获取处理记录详情
        
        Args:
            process_id: 处理记录ID
            user_id: 用户ID（可选，用于权限验证）
            
        Returns:
            Dict: 处理记录详情，包含关联的图片信息
        """
        process_record = self.db_ops.get_process_record_by_id(process_id)
        if not process_record:
            return None
        
        # 权限验证
        if user_id and process_record["user_id"] != user_id:
            return None
        
        # 获取关联的图片信息
        face_image = self.db_ops.get_image_record_by_id(process_record["face_image_id"])
        cloth_image = self.db_ops.get_image_record_by_id(process_record["cloth_image_id"])
        result_image = self.db_ops.get_image_record_by_id(process_record["result_image_id"])
        
        return {
            "process_record": process_record,
            "face_image": face_image,
            "cloth_image": cloth_image,
            "result_image": result_image
        }
    
    def approve_process_record(self, process_id: int) -> bool:
        """
        审核通过处理记录
        
        Args:
            process_id: 处理记录ID
            
        Returns:
            bool: 操作是否成功
        """
        try:
            # 检查记录是否存在
            process_record = self.db_ops.get_process_record_by_id(process_id)
            if not process_record:
                logger.error(f"处理记录不存在: {process_id}")
                return False
            
            # 更新状态为已审核
            update_data = {
                "status": "已审核"
            }
            
            updated_record = self.db_ops.update_process_record(process_id, update_data)
            
            if updated_record:
                logger.info(f"处理记录 {process_id} 审核通过")
                return True
            else:
                logger.error(f"更新处理记录状态失败: {process_id}")
                return False
                
        except Exception as e:
            logger.error(f"审核处理记录异常: {str(e)}", exc_info=True)
            return False

    def delete_result_image(self, process_id: int) -> bool:
        """
        删除处理记录的结果图片：
        - 清空 ProcessRecord.result_image_id 与 generated_text 中与图片相关的内容保持不变
        - 可选地将状态置为“待审核”或维持原状态；此处不更改状态
        - 不删除底层图片文件与 ImageRecord，仅解除关联，避免误删数据
        """
        try:
            record = self.db_ops.get_process_record_by_id(process_id)
            if not record:
                logger.error(f"处理记录不存在: {process_id}")
                return False
            update_data = {
                "result_image_id": None
            }
            updated = self.db_ops.update_process_record(process_id, update_data)
            if not updated:
                logger.error(f"更新处理记录失败: {process_id}")
                return False
            logger.info(f"处理记录 {process_id} 已解除结果图片关联")
            return True
        except Exception as e:
            logger.error(f"删除结果图片失败: {str(e)}", exc_info=True)
            return False


# 创建全局服务实例
ai_swap_cloth_service = AISwapClothService()

def process_swap_cloth_with_record(
    user_id: int,
    raw_image_id: int,
    cloth_image_id: int,
    **kwargs
) -> Dict[str, Any]:
    """
    便捷函数：执行完整的换衣服流程

    Args:
        user_id: 用户ID
        raw_image_id: 原始图片ID
        cloth_image_id: 衣服图片ID
        **kwargs: 其他可选参数

    Returns:
        Dict: 包含处理结果的字典
    """
    return ai_swap_cloth_service.process_swap_cloth_with_record(
        user_id, raw_image_id, cloth_image_id, **kwargs
    )


if __name__ == "__main__":
    try:
        # 假设用户ID为1，原始图片ID为1，衣服图片ID为1
        result = process_swap_cloth_with_record(
            user_id=5,
            raw_image_id=201,
            cloth_image_id=404
        )
        
        if result["success"]:
            print("处理成功！")
            print(f"处理记录ID: {result['process_record_id']}")
            print(f"结果图片ID: {result['result_image_id']}")
            print(f"文案: {result['copywriter_text']}")
        else:
            print(f"处理失败: {result['error']}")
            
    except Exception as e:
        print(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()