Tolan/backend/services/message_service.py

"""
消息存储服务（SQLAlchemy ORM 版本）
"""
import uuid
import json
import asyncio
from datetime import datetime
from typing import Awaitable, Callable, Dict, Optional
import logging
from sqlalchemy import select, update, func
from sqlalchemy.orm import selectinload
from database import db_manager
from models_db import Message, User
import redis.asyncio as aioredis

logger = logging.getLogger(__name__)


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

        Args:
            db_manager: 数据库管理器实例（可选，用于延迟初始化）
        """
        self.db = db_manager
        self.redis = None
        self._redis_provider: Optional[Callable[[], Awaitable[aioredis.Redis]]] = None
        if db_manager:
            # 🔥 延迟获取Redis（在首次使用时）
            self._redis_provider = db_manager.get_redis

    async def set_db_manager(self, db_manager):
        """设置数据库管理器（用于延迟初始化）"""
        self.db = db_manager
        self._redis_provider = db_manager.get_redis
        self.redis = None
    
    async def _ensure_redis(self):
        """确保Redis连接可用"""
        if not self._redis_provider:
            if not self.db:
                raise RuntimeError("Redis 未初始化，缺少数据库管理器")
            self._redis_provider = self.db.get_redis
        redis_client = await self._redis_provider()
        if not redis_client:
            raise RuntimeError("Redis 客户端不可用")
        self.redis = redis_client
        return redis_client
    
    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:
            # 🔥 确保Redis连接可用
            await self._ensure_redis()
            
            # 1. 获取当前轮次
            logger.debug(f"📥 [保存消息] 步骤1: 获取当前轮次")
            current_round = await self._get_or_init_round(user_id, role)
            
            # 2. 生成消息ID
            logger.debug(f"📥 [保存消息] 步骤2: 生成消息ID")
            message_id = str(uuid.uuid4())
            logger.debug(f"📥 [保存消息] 消息ID: {message_id}")
            
            # 3. 保存到MySQL（使用 ORM）
            logger.debug(f"📥 [保存消息] 步骤3: 保存到MySQL")
            async with db_manager.get_session() as session:
                message = Message(
                    message_id=message_id,
                    user_id=user_id,
                    role=role,
                    content=content,
                    round_number=current_round,
                    emotion=emotion,
                    message_metadata=metadata,
                    token_count=len(content) // 2  # 粗略估算
                )
                session.add(message)
                await session.commit()
            
            # 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)
                
                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}")
            
            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获取（增加超时控制）
            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（新轮次）
                    max_round = await asyncio.wait_for(
                        self._get_max_round(user_id),
                        timeout=3.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._get_latest_user_round(user_id),
                            timeout=3.0
                        )
                        if result:
                            current_round = result
                            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}")
                max_round = await asyncio.wait_for(
                    self._get_max_round(user_id),
                    timeout=3.0  # 3秒超时
                )
                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}")
            
            return current_round
            
        except asyncio.TimeoutError:
            logger.warning(f"⚠️ Redis/MySQL超时，使用降级方案: user={user_id}")
            # 降级方案：尝试快速从MySQL查询，失败则返回1
            try:
                max_round = await asyncio.wait_for(
                    self._get_max_round(user_id),
                    timeout=2.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 _get_max_round(self, user_id: str) -> int:
        """使用 ORM 获取用户最大轮次"""
        try:
            async with db_manager.get_session() as session:
                result = await session.execute(
                    select(func.max(Message.round_number))
                    .where(Message.user_id == user_id)
                )
                max_round = result.scalar()
                return max_round if max_round else 0
        except Exception as e:
            logger.error(f"❌ 获取最大轮次失败: user={user_id}, error={e}")
            return 0
    
    async def _get_latest_user_round(self, user_id: str) -> Optional[int]:
        """使用 ORM 获取最近一条 user 消息的轮次"""
        try:
            async with db_manager.get_session() as session:
                result = await session.execute(
                    select(Message.round_number)
                    .where(Message.user_id == user_id)
                    .where(Message.role == 'user')
                    .order_by(Message.created_at.desc())
                    .limit(1)
                )
                round_number = result.scalar_one_or_none()
                return round_number
        except Exception as e:
            logger.error(f"❌ 获取最近user消息轮次失败: user={user_id}, error={e}")
            return None
    
    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中的用户统计（使用 ORM）"""
        try:
            async with db_manager.get_session() as session:
                await session.execute(
                    update(User)
                    .where(User.user_id == user_id)
                    .values(
                        total_messages=User.total_messages + 1,
                        total_rounds=round_number,
                        last_active_at=func.now()
                    )
                )
                await session.commit()
        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:
            redis = await self._ensure_redis()
            meta_key = f"user:{user_id}:meta"
            current_round = await redis.hget(meta_key, 'current_round')
            
            if current_round:
                return int(current_round)
            
            # 从MySQL查询（使用 ORM）
            max_round = await self._get_max_round(user_id)
            return max_round if max_round else 0
            
        except Exception as e:
            logger.error(f"❌ 获取当前轮次失败: user={user_id}, error={e}")
            return 0


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