ForcePilot/backend/package/yuxi/channel/extensions/rocketchat/monitor.py

from __future__ import annotations

import logging
import time

from yuxi.channel.extensions.rocketchat.client import RocketChatClient
from yuxi.channel.extensions.rocketchat.dedup import ClaimableDedupe
from yuxi.channel.extensions.rocketchat.errors import RocketChatError
from yuxi.channel.extensions.rocketchat.format import normalize_message
from yuxi.channel.extensions.rocketchat.gating import RocketChatGating
from yuxi.channel.extensions.rocketchat.outbound import map_room_type_to_chat_type
from yuxi.channel.extensions.rocketchat.security import RocketChatSecurityAdapter
from yuxi.channel.extensions.rocketchat.threading import RocketChatThreadingAdapter

logger = logging.getLogger(__name__)


class RocketChatMonitor:
    def __init__(
        self,
        client: RocketChatClient,
        account_id: str,
        account: dict,
        on_message=None,
        on_dispatched=None,
    ):
        self.client = client
        self.account_id = account_id
        self.account = account
        self._on_message = on_message
        self._on_dispatched = on_dispatched

        self.dedupe = ClaimableDedupe(ttl_ms=300_000, max_size=2000)
        self.security = RocketChatSecurityAdapter()
        self.gating = RocketChatGating(account)
        self.threading = RocketChatThreadingAdapter(account)

        self.bot_user_id: str = ""
        self.bot_username: str = ""

    async def handle_message(self, raw_msg: dict) -> dict | None:
        try:
            msg_id = raw_msg.get("_id", "")
            room_id = raw_msg.get("rid", "")
            user = raw_msg.get("u", {})
            user_id = user.get("_id", "")
            username = user.get("username", "")

            if not msg_id or not room_id:
                return None

            if not user_id:
                logger.debug("Message missing user_id, dropping")
                return None

            if user_id == self.bot_user_id:
                return None

            msg_type = raw_msg.get("t", "")
            if msg_type:
                logger.debug("Skipping system message type: %s", msg_type)
                return None

            dedupe_key = f"{self.account_id}:{msg_id}"
            result = self.dedupe.claim(dedupe_key)
            if result == "duplicate":
                logger.debug("Duplicate message %s, dropping", msg_id)
                return None

            try:
                room_info = await self.client.fetch_room_info(room_id)
                room_t = room_info.get("room", {}).get("t", "c")
                chat_type = map_room_type_to_chat_type(room_t)
            except RocketChatError:
                chat_type = "channel"

            allowed, reason = self.security.check_sender_access(
                chat_type,
                user_id,
                self.account,
                username,
            )
            if not allowed:
                logger.info(
                    "Access denied for %s in %s: %s",
                    user_id,
                    room_id,
                    reason,
                )
                return {
                    "status": "blocked",
                    "reason": reason,
                    "room_id": room_id,
                    "chat_type": chat_type,
                }

            raw_text = raw_msg.get("msg", "")
            normalized = normalize_message(raw_text, self.bot_username)

            was_mentioned = self.gating.was_mentioned(
                raw_text,
                self.bot_user_id,
                self.bot_username,
            )

            should_respond, gate_reason = self.gating.should_respond(
                chat_type,
                normalized,
                was_mentioned,
            )

            if not should_respond:
                logger.debug("Gate blocked: %s", gate_reason)
                return {
                    "status": "gated",
                    "reason": gate_reason,
                    "room_id": room_id,
                    "chat_type": chat_type,
                }

            thread_id = raw_msg.get("tmid", "")

            thread_context = self.threading.build_thread_context(
                thread_id,
                msg_id,
                room_id,
                chat_type,
            )

            ts_field = raw_msg.get("ts", {})
            if isinstance(ts_field, dict):
                timestamp = ts_field.get("$date", time.time() * 1000) / 1000.0
            else:
                timestamp = time.time()

            file_id = raw_msg.get("file", {}).get("_id", "")
            mentions = [m.get("_id", "") for m in raw_msg.get("mentions", [])]

            unified = {
                "channel_type": "rocketchat",
                "account_id": self.account_id,
                "msg_id": msg_id,
                "room_id": room_id,
                "chat_type": chat_type,
                "sender_id": user_id,
                "sender_name": username,
                "content": normalized,
                "raw_content": raw_text,
                "thread_id": thread_id,
                "thread_context": thread_context,
                "file_ids": [file_id] if file_id else [],
                "mentions": mentions,
                "timestamp": timestamp,
            }

            if self._on_message:
                try:
                    await self._on_message(unified)
                except Exception as e:
                    logger.error("Error in on_message handler: %s", e)

            return {
                "status": "processed",
                "room_id": room_id,
                "chat_type": chat_type,
                "msg_id": msg_id,
                "unified": unified,
            }

        except Exception as e:
            logger.error("Error handling message %s: %s", raw_msg.get("_id", "?"), e)
            return None

    def reset(self) -> None:
        self.dedupe.reset()