ai-chat-system/services/matrix_service.py

"""
Matrix Bot 服务 - 处理Matrix消息收发
"""
import asyncio
import json
from typing import Optional, Callable
from nio import AsyncClient, MatrixRoom, RoomMessageText, LoginResponse, SyncResponse
import logging

from models import SessionLocal, User, Conversation, Message, MatrixRoomMapping, SystemConfig
from services.conversation_service import ConversationService

logger = logging.getLogger(__name__)


class MatrixBot:
    def __init__(self):
        self.client: Optional[AsyncClient] = None
        self.homeserver: str = ""
        self.username: str = ""
        self.password: str = ""
        self.access_token: str = ""
        self.is_running: bool = False
        self.on_message_callback: Optional[Callable] = None
    
    async def init_from_config(self):
        """从数据库配置初始化"""
        db = SessionLocal()
        try:
            configs = {c.key: c.value for c in db.query(SystemConfig).all()}
            self.homeserver = configs.get('matrix_homeserver', 'https://matrix.tphai.com')
            self.username = configs.get('matrix_username', '')
            self.password = configs.get('matrix_password', '')
            self.access_token = configs.get('matrix_access_token', '')
            
            if self.username and (self.password or self.access_token):
                await self.connect()
        finally:
            db.close()
    
    async def connect(self):
        """连接到Matrix服务器"""
        if not self.homeserver or not self.username:
            logger.warning("Matrix配置不完整，跳过连接")
            return False
        
        try:
            # 创建客户端
            self.client = AsyncClient(
                self.homeserver, 
                self.username,
                store_path="/tmp/matrix_store"
            )
            
            # 如果有access_token，直接设置
            if self.access_token:
                self.client.access_token = self.access_token
                self.client.user_id = self.username
                logger.info(f"Matrix已设置access_token: {self.username}")
                self.is_running = True
                return True
            
            # 否则使用密码登录
            response = await self.client.login(self.password)
            
            if isinstance(response, LoginResponse):
                logger.info(f"Matrix连接成功: {self.username}")
                self.is_running = True
                return True
            else:
                logger.error(f"Matrix登录失败: {response}")
                return False
        except Exception as e:
            logger.error(f"Matrix连接错误: {e}")
            return False
    
    async def start_sync(self, message_handler: Callable = None):
        """开始同步消息"""
        if not self.client:
            logger.warning("Matrix客户端未初始化")
            return
        
        self.on_message_callback = message_handler
        
        # 注册消息处理器
        self.client.add_event_callback(self._handle_room_message, RoomMessageText)
        
        # 首先执行一次同步获取房间信息
        try:
            sync_response = await self.client.sync(timeout=10000)
            logger.info(f"Matrix初始同步完成")
        except Exception as e:
            logger.warning(f"Matrix初始同步失败: {e}, 将尝试继续")
        
        # 启动后台同步任务（不阻塞）
        asyncio.create_task(self._sync_loop())
        logger.info("Matrix同步任务已启动")
    
    async def _sync_loop(self):
        """后台同步循环"""
        while self.is_running:
            try:
                # 使用较短的超时，避免长时间阻塞
                await self.client.sync(timeout=5000)
                await asyncio.sleep(1)  # 每秒同步一次
            except Exception as e:
                logger.warning(f"Matrix同步错误: {e}")
                await asyncio.sleep(5)  # 出错后等待5秒再重试
    
    async def _handle_room_message(self, room: MatrixRoom, event: RoomMessageText):
        """处理收到的房间消息"""
        # 忽略自己发送的消息
        if event.sender == self.username:
            return
        
        db = SessionLocal()
        try:
            conv_service = ConversationService(db)
            
            # 获取或创建Matrix用户
            matrix_user_id = event.sender
            user = conv_service.get_or_create_user(
                user_id=matrix_user_id,
                display_name=room.users.get(matrix_user_id, {}).get('display_name', matrix_user_id),
                user_type='matrix',
                matrix_user_id=matrix_user_id
            )
            
            # 获取或创建房间映射
            mapping = db.query(MatrixRoomMapping).filter(
                MatrixRoomMapping.room_id == room.room_id
            ).first()
            
            if not mapping:
                # 为这个房间创建新会话
                conversation = conv_service.create_conversation(user.id, title=f"Matrix: {room.display_name}")
                mapping = MatrixRoomMapping(
                    room_id=room.room_id,
                    user_id=user.id,
                    conversation_id=conversation.id
                )
                db.add(mapping)
                db.commit()
                db.refresh(mapping)
            
            # 保存用户消息
            conv_service.add_message(
                conversation_id=mapping.conversation_id,
                role='user',
                content=event.body,
                source='matrix',
                extra_data={'event_id': event.event_id, 'room_id': room.room_id}
            )
            
            # 调用外部消息处理器
            if self.on_message_callback:
                await self.on_message_callback(
                    conversation_id=mapping.conversation.conversation_id,
                    user_message=event.body,
                    user_id=user.user_id,
                    room_id=room.room_id
                )
        finally:
            db.close()
    
    async def send_message(self, room_id: str, message: str):
        """发送消息到Matrix房间"""
        if not self.client:
            return False
        
        try:
            await self.client.room_send(
                room_id=room_id,
                message_type="m.room.message",
                content={
                    "msgtype": "m.text",
                    "body": message
                }
            )
            return True
        except Exception as e:
            logger.error(f"发送Matrix消息失败: {e}")
            return False
    
    async def get_room_id_for_user(self, user_matrix_id: str) -> Optional[str]:
        """获取与用户的对话房间ID"""
        db = SessionLocal()
        try:
            user = db.query(User).filter(User.matrix_user_id == user_matrix_id).first()
            if user:
                mapping = db.query(MatrixRoomMapping).filter(
                    MatrixRoomMapping.user_id == user.id
                ).first()
                if mapping:
                    return mapping.room_id
            return None
        finally:
            db.close()
    
    async def disconnect(self):
        """断开连接"""
        if self.client:
            await self.client.logout()
            await self.client.close()
            self.is_running = False


# 全局实例
matrix_bot = MatrixBot()