Tolan/backend/services/message_service.py

"""
消息存储服务
"""
import uuid
import json
import asyncio
from datetime import datetime
from typing import Dict, Optional
import logging

logger = logging.getLogger(__name__)


class MessageService:
    """消息存储服务"""
    
    def __init__(self, db_manager=None):
        """
        初始化消息服务

        Args:
            db_manager: 数据库管理器实例（可选，用于延迟初始化）
        """
        self.db = db_manager
        self.redis = None
        if db_manager:
            self.redis = db_manager.get_redis()

    def set_db_manager(self, db_manager):
        """设置数据库管理器（用于延迟初始化）"""
        self.db = db_manager
        self.redis = db_manager.get_redis()
    
    async def save_message(
        self,
        user_id: str,
        role: str,
        content: str,
        emotion: Optional[str] = None,
        metadata: Optional[dict] = None
    ) -> Dict:
        """
        保存消息到MySQL和Redis
        
        Args:
            user_id: 用户ID
            role: 角色 ('user', 'assistant', 'system')
            content: 消息内容
            emotion: 情感标签（可选）
            metadata: 扩展信息（可选）
        
        Returns:
            消息信息（包含round_number）
        """
        
        try:
            logger.info("🚨 [TEST] save_message 方法被调用了！")
            logger.info(f"📥 [保存消息] 开始: user={user_id}, role={role}, content_length={len(content)}")
            
            # 1. 获取当前轮次
            logger.debug(f"📥 [保存消息] 步骤1: 获取当前轮次")
            current_round = await self._get_or_init_round(user_id, role)
            logger.info(f"📊 [保存消息] user={user_id}, role={role}, current_round={current_round}")
            
            # 2. 生成消息ID
            logger.debug(f"📥 [保存消息] 步骤2: 生成消息ID")
            message_id = str(uuid.uuid4())
            logger.debug(f"📥 [保存消息] 消息ID: {message_id}")
            
            # 3. 保存到MySQL
            logger.debug(f"📥 [保存消息] 步骤3: 保存到MySQL")
            query = """
                INSERT INTO messages 
                (message_id, user_id, role, content, round_number, emotion, metadata, token_count)
                VALUES (%(message_id)s, %(user_id)s, %(role)s, %(content)s, %(round_number)s, 
                        %(emotion)s, %(metadata)s, %(token_count)s)
            """
            
            await self.db.execute_update(query, {
                "message_id": message_id,
                "user_id": user_id,
                "role": role,
                "content": content,
                "round_number": current_round,
                "emotion": emotion,
                "metadata": json.dumps(metadata) if metadata else None,
                "token_count": len(content) // 2  # 粗略估算
            })
            
            logger.info(f"✅ [保存消息] MySQL保存成功: [Round {current_round}] user={user_id}, role={role}, message_id={message_id}")
            
            # 4. 构建消息对象
            logger.debug(f"📥 [保存消息] 步骤4: 构建消息对象")
            message_obj = {
                "message_id": message_id,
                "role": role,
                "content": content,
                "round_number": current_round,
                "emotion": emotion,
                "created_at": datetime.now().isoformat()
            }
            
            # 5. 更新Redis缓存（最近15轮=30条消息）
            logger.debug(f"📥 [保存消息] 步骤5: 更新Redis缓存")
            await self._cache_message(user_id, message_obj)
            logger.debug(f"📥 [保存消息] Redis缓存更新完成")
            
            # 6. 更新用户统计
            logger.debug(f"📥 [保存消息] 步骤6: 更新用户统计")
            await self._update_user_stats(user_id, current_round, role)
            logger.debug(f"📥 [保存消息] 用户统计更新完成")
            
            # 7. 检查是否需要触发总结（每15轮）
            logger.debug(f"📥 [保存消息] 步骤7: 检查是否需要触发总结")
            if role == 'assistant':
                is_multiple_of_15 = (current_round % 15 == 0)
                remainder = current_round % 15
                logger.info(
                    f"🔍 [总结触发检查] user={user_id}, role={role}, "
                    f"current_round={current_round}, "
                    f"是否15的倍数={is_multiple_of_15}, "
                    f"余数={remainder}, "
                    f"计算: {current_round} % 15 = {remainder}"
                )
                
                if is_multiple_of_15:
                    # 异步触发总结任务（使用 create_task 并添加异常处理）
                    logger.info(f"🔄 [总结触发] 触发总结任务: user={user_id}, round={current_round}")
                    task = asyncio.create_task(
                        self._trigger_summary(user_id, current_round)
                    )
                    logger.debug(f"🔄 [总结触发] 总结任务已创建: user={user_id}, task={task}")
                    
                    # 添加异常回调，确保任务失败时能记录日志
                    def log_task_exception(task):
                        try:
                            task.result()  # 这会抛出异常（如果有）
                        except Exception as e:
                            logger.error(f"❌ [总结任务] 执行失败: user={user_id}, error={e}")
                            import traceback
                            logger.error(f"❌ [总结任务] 异常堆栈: {traceback.format_exc()}")
                    task.add_done_callback(log_task_exception)
                    logger.debug(f"🔄 [总结触发] 异常回调已添加: user={user_id}")
                else:
                    logger.debug(f"🔍 [总结触发检查] 不需要触发总结: user={user_id}, current_round={current_round} 不是15的倍数")
            else:
                logger.debug(f"🔍 [总结触发检查] 跳过（非assistant消息）: user={user_id}, role={role}")
            
            logger.info(f"✅ [保存消息] 完成: user={user_id}, role={role}, round={current_round}, message_id={message_id}")
            return message_obj
            
        except Exception as e:
            logger.error(f"❌ 保存消息失败: user={user_id}, error={e}")
            raise
    
    async def _get_or_init_round(self, user_id: str, role: str) -> int:
        """
        获取当前轮次
        规则：user消息开始新轮次，assistant消息沿用当前轮次
        
        Args:
            user_id: 用户ID
            role: 角色
        
        Returns:
            当前轮次
        """
        
        meta_key = f"user:{user_id}:meta"
        current_round = 1  # 默认值
        
        logger.debug(f"🔍 [获取轮次] 开始: user={user_id}, role={role}, meta_key={meta_key}")
        
        try:
            # 从Redis获取（增加超时控制）
            import asyncio
            logger.debug(f"🔍 [获取轮次] 从Redis读取: user={user_id}, key={meta_key}")
            current_round_str = await asyncio.wait_for(
                self.redis.hget(meta_key, 'current_round'),
                timeout=2.0  # 2秒超时
            )
            logger.debug(f"🔍 [获取轮次] Redis返回值: user={user_id}, current_round_str={current_round_str}")
            
            if current_round_str:
                redis_round = int(current_round_str)
                logger.debug(f"🔍 [获取轮次] Redis中的轮次: user={user_id}, redis_round={redis_round}")
                
                if role == 'user':
                    # user消息：从MySQL查询最大轮次，然后+1（新轮次）
                    result = await asyncio.wait_for(
                        self.db.execute_one(
                            "SELECT MAX(round_number) as max_round FROM messages WHERE user_id = %(user_id)s",
                            {"user_id": user_id}
                        ),
                        timeout=3.0
                    )
                    max_round = result['max_round'] if result and result['max_round'] else 0
                    current_round = max_round + 1
                    logger.debug(f"🔍 [获取轮次] user消息，新轮次: user={user_id}, max_round={max_round} -> current_round={current_round}")
                else:
                    # assistant消息：使用Redis中的当前轮次（应该和对应的user消息相同）
                    # 如果Redis中没有值或值异常，则查询MySQL中最近一条user消息的round_number
                    if redis_round > 0:
                        current_round = redis_round
                        logger.debug(f"🔍 [获取轮次] assistant消息，使用Redis轮次: user={user_id}, current_round={current_round}")
                    else:
                        # Redis值异常，查询MySQL中最近一条user消息的round_number
                        result = await asyncio.wait_for(
                            self.db.execute_one(
                                "SELECT round_number FROM messages WHERE user_id = %(user_id)s AND role = 'user' ORDER BY created_at DESC LIMIT 1",
                                {"user_id": user_id}
                            ),
                            timeout=3.0
                        )
                        if result and result.get('round_number'):
                            current_round = result['round_number']
                            logger.debug(f"🔍 [获取轮次] assistant消息，使用MySQL最近user消息轮次: user={user_id}, current_round={current_round}")
                        else:
                            # 没有user消息，说明这是第一轮对话
                            current_round = 1
                            logger.debug(f"🔍 [获取轮次] assistant消息，没有user消息，使用默认轮次: user={user_id}, current_round={current_round}")
            else:
                # 从MySQL获取最大轮次
                logger.debug(f"🔍 [获取轮次] Redis无值，从MySQL查询: user={user_id}")
                result = await asyncio.wait_for(
                    self.db.execute_one(
                        "SELECT MAX(round_number) as max_round FROM messages WHERE user_id = %(user_id)s",
                        {"user_id": user_id}
                    ),
                    timeout=3.0  # 3秒超时
                )
                
                max_round = result['max_round'] if result and result['max_round'] else 0
                logger.debug(f"🔍 [获取轮次] MySQL最大轮次: user={user_id}, max_round={max_round}")
                
                if role == 'user':
                    current_round = max_round + 1
                    logger.debug(f"🔍 [获取轮次] user消息，新轮次: user={user_id}, current_round={current_round}")
                else:
                    # assistant 消息：总是使用 max_round + 1，因为每一轮对话都有唯一的轮次号
                    current_round = max_round + 1
                    logger.debug(f"🔍 [获取轮次] assistant消息，新轮次: user={user_id}, current_round={current_round}")
            
            # 确保 current_round 至少为 1（防止为 0 的情况）
            if current_round <= 0:
                logger.warning(f"⚠️ 计算出的轮次异常: user={user_id}, role={role}, current_round={current_round}, 使用默认值1")
                current_round = 1
            
            # 更新Redis（user消息时更新，assistant消息时如果计算出的值更大也更新）
            should_update_redis = False
            if role == 'user':
                should_update_redis = True
                logger.debug(f"🔍 [获取轮次] user消息，需要更新Redis: user={user_id}, current_round={current_round}")
            elif role == 'assistant':
                # assistant 消息：如果计算出的 current_round 大于 Redis 中的值，也更新
                # 这样可以修复 Redis 中异常值（如 0）的问题
                try:
                    logger.debug(f"🔍 [获取轮次] assistant消息，检查是否需要更新Redis: user={user_id}")
                    redis_round_str = await asyncio.wait_for(
                        self.redis.hget(meta_key, 'current_round'),
                        timeout=1.0
                    )
                    if redis_round_str:
                        redis_round = int(redis_round_str)
                        logger.debug(f"🔍 [获取轮次] Redis中的轮次: user={user_id}, redis_round={redis_round}, current_round={current_round}")
                        if current_round > redis_round:
                            should_update_redis = True
                            logger.info(f"🔄 修复Redis轮次: user={user_id}, Redis={redis_round} -> {current_round}")
                        else:
                            logger.debug(f"🔍 [获取轮次] 不需要更新Redis: user={user_id}, redis_round={redis_round} >= current_round={current_round}")
                    else:
                        # Redis 中没有值，也更新
                        should_update_redis = True
                        logger.debug(f"🔍 [获取轮次] Redis无值，需要更新: user={user_id}, current_round={current_round}")
                except Exception as e:
                    logger.warning(f"⚠️ Redis读取失败（不影响主流程）: user={user_id}, error={e}")
            
            if should_update_redis:
                try:
                    logger.debug(f"🔍 [获取轮次] 更新Redis: user={user_id}, current_round={current_round}")
                    # 更新 user:{user_id}:meta 中的 current_round
                    await asyncio.wait_for(
                        self.redis.hset(meta_key, 'current_round', current_round),
                        timeout=1.0
                    )
                    await asyncio.wait_for(
                        self.redis.expire(meta_key, 86400),
                        timeout=1.0
                    )
                    # 同时更新 round:{user_id}（保持一致性）
                    round_key = f"round:{user_id}"
                    await asyncio.wait_for(
                        self.redis.set(round_key, current_round),
                        timeout=1.0
                    )
                    await asyncio.wait_for(
                        self.redis.expire(round_key, 86400),
                        timeout=1.0
                    )
                    logger.debug(f"🔍 [获取轮次] Redis更新成功: user={user_id}, current_round={current_round}")
                except Exception as e:
                    logger.warning(f"⚠️ Redis写入失败（不影响主流程）: user={user_id}, error={e}")
            
            logger.info(f"✅ [获取轮次] 完成: user={user_id}, role={role}, current_round={current_round}")
            logger.info(f"🎯 [轮次计算] user={user_id}, role={role}, 计算结果: current_round={current_round}, 下一个15的倍数={(current_round // 15 + 1) * 15}")
            return current_round
            
        except asyncio.TimeoutError:
            logger.warning(f"⚠️ Redis/MySQL超时，使用降级方案: user={user_id}")
            # 降级方案：尝试快速从MySQL查询，失败则返回1
            try:
                result = await asyncio.wait_for(
                    self.db.execute_one(
                        "SELECT MAX(round_number) as max_round FROM messages WHERE user_id = %(user_id)s",
                        {"user_id": user_id}
                    ),
                    timeout=2.0
                )
                max_round = result['max_round'] if result and result['max_round'] else 0
                if role == 'user':
                    current_round = max_round + 1
                else:
                    # assistant 消息：总是使用 max_round + 1，因为每一轮对话都有唯一的轮次号
                    current_round = max_round + 1
                # 确保至少为 1
                return max(current_round, 1)
            except:
                logger.error(f"❌ 降级方案也失败，返回默认轮次1")
                return 1
            
        except Exception as e:
            logger.error(f"❌ 获取轮次失败: user={user_id}, error={e}")
            # 最终降级：返回1
            return 1
    
    async def _cache_message(self, user_id: str, message: Dict):
        """
        缓存消息到Redis（仅保留最近30条）
        
        Args:
            user_id: 用户ID
            message: 消息对象
        """
        
        try:
            cache_key = f"user:{user_id}:recent_messages"
            logger.debug(f"🔍 [缓存消息] 开始: user={user_id}, cache_key={cache_key}, round={message.get('round_number')}")
            await self.redis.lpush(cache_key, json.dumps(message, ensure_ascii=False))
            await self.redis.ltrim(cache_key, 0, 29)  # 保留最近30条
            await self.redis.expire(cache_key, 86400)  # 24小时过期
            logger.debug(f"✅ [缓存消息] 完成: user={user_id}, round={message.get('round_number')}")
        except Exception as e:
            logger.warning(f"⚠️ [缓存消息] 失败: user={user_id}, error={e}")
    
    async def _update_user_stats(self, user_id: str, round_number: int, role: str):
        """
        更新用户统计信息
        
        Args:
            user_id: 用户ID
            round_number: 轮次
            role: 角色
        """
        
        try:
            logger.debug(f"🔍 [更新统计] 开始: user={user_id}, round={round_number}, role={role}")
            
            # 更新Redis元数据
            meta_key = f"user:{user_id}:meta"
            await self.redis.hincrby(meta_key, 'total_messages', 1)
            logger.debug(f"🔍 [更新统计] Redis消息数+1: user={user_id}")
            
            if role == 'assistant':
                await self.redis.hset(meta_key, 'total_rounds', round_number)
                logger.debug(f"🔍 [更新统计] Redis轮次更新: user={user_id}, round={round_number}")
            
            await self.redis.expire(meta_key, 86400)
            
            # 更新MySQL（异步，不阻塞）
            logger.debug(f"🔍 [更新统计] 创建MySQL更新任务: user={user_id}")
            asyncio.create_task(
                self._update_user_stats_db(user_id, round_number)
            )
            
            logger.debug(f"✅ [更新统计] 完成: user={user_id}, round={round_number}")
        except Exception as e:
            logger.warning(f"⚠️ [更新统计] 失败: user={user_id}, error={e}")
    
    async def _update_user_stats_db(self, user_id: str, round_number: int):
        """更新MySQL中的用户统计"""
        try:
            query = """
                UPDATE users 
                SET total_messages = total_messages + 1,
                    total_rounds = %(round_number)s,
                    last_active_at = NOW()
                WHERE user_id = %(user_id)s
            """
            await self.db.execute_update(query, {
                "user_id": user_id,
                "round_number": round_number
            })
        except Exception as e:
            logger.error(f"❌ 更新MySQL用户统计失败: user={user_id}, error={e}")
    
    async def _trigger_summary(self, user_id: str, current_round: int):
        """
        触发异步总结任务
        
        Args:
            user_id: 用户ID
            current_round: 当前轮次
        """
        try:
            logger.info(f"🚀 [总结任务] ========== 开始执行 ==========")
            logger.info(f"🚀 [总结任务] user={user_id}, current_round={current_round}")
            
            # 参数验证
            if current_round <= 0:
                logger.warning(f"⚠️ [总结任务] 无效的当前轮次: user={user_id}, current_round={current_round}")
                return
            
            logger.debug(f"🔍 [总结任务] 步骤1: 导入SummaryService")
            from .summary_service import SummaryService
            
            logger.debug(f"🔍 [总结任务] 步骤2: 创建SummaryService实例")
            summary_service = SummaryService(self.db)
            
            # 计算总结范围
            # 当前轮次是15的倍数，总结1到current_round
            start_round = 1
            end_round = current_round
            
            logger.info(f"📝 [总结任务] 准备生成总结: user={user_id}, start_round={start_round}, end_round={end_round}")
            logger.debug(f"🔍 [总结任务] 步骤3: 调用generate_summary")
            
            await summary_service.generate_summary(
                user_id=user_id,
                start_round=start_round,
                end_round=end_round
            )
            
            logger.info(f"✅ [总结任务] ========== 完成 ==========")
            logger.info(f"✅ [总结任务] user={user_id}, current_round={current_round}")
            
        except Exception as e:
            logger.error(f"❌ [总结任务] ========== 执行失败 ==========")
            logger.error(f"❌ [总结任务] user={user_id}, current_round={current_round}, error={e}")
            import traceback
            logger.error(f"❌ [总结任务] 异常堆栈:\n{traceback.format_exc()}")
    
    async def get_user_current_round(self, user_id: str) -> int:
        """
        获取用户当前对话轮次
        
        Args:
            user_id: 用户ID
        
        Returns:
            当前轮次（从1开始）
        """
        try:
            meta_key = f"user:{user_id}:meta"
            current_round = await self.redis.hget(meta_key, 'current_round')
            
            if current_round:
                return int(current_round)
            
            # 从MySQL查询
            result = await self.db.execute_one(
                "SELECT MAX(round_number) as max_round FROM messages WHERE user_id = %(user_id)s",
                {"user_id": user_id}
            )
            
            return result['max_round'] if result and result['max_round'] else 0
            
        except Exception as e:
            logger.error(f"❌ 获取当前轮次失败: user={user_id}, error={e}")
            return 0


# 全局消息服务实例
message_service = MessageService()