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)
            
            # 如果有access_token，直接使用
            if self.access_token:
                self.client.access_token = self.access_token
                logger.info(f"Matrix连接成功(使用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:
            return
        
        self.on_message_callback = message_handler
        
        # 注册消息处理器
        self.client.add_event_callback(self._handle_room_message, RoomMessageText)
        
        # 开始同步循环
        await self.client.sync_forever(timeout=30000)
    
    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()