ForcePilot/backend/package/yuxi/channels/adapters/matrix/normalizer.py

from __future__ import annotations

import re
from datetime import datetime, UTC
from typing import TYPE_CHECKING, Any

from nio import (
    RoomMessageAudio,
    RoomMessageFile,
    RoomMessageFormatted,
    RoomMessageImage,
    RoomMessageText,
    RoomMessageVideo,
)

from yuxi.channels.models import (
    Attachment,
    ChannelIdentity,
    ChannelMessage,
    ChannelType,
    ChatType,
    EventType,
    MentionsInfo,
    MessageType,
)

from .rooms_config import get_room_require_mention
from .thread_utils import is_thread_event
from .location import parse_location_event
from .sticker import is_sticker_event, parse_sticker

if TYPE_CHECKING:
    from yuxi.channels.adapters.matrix.adapter import MatrixAdapter

_MXID_PATTERN = re.compile(r'data-mx-username="@([^"]+)"')
_MXID_URL_PATTERN = re.compile(r'href="https://matrix\.to/#/(@[^"]+)"')
_PLAIN_MXID_PATTERN = re.compile(r"@([\w._=\-/]+):([\w.\-]+(?:\.[\w.\-]+)*)")


class MatrixNormalizer:
    def __init__(self, adapter: MatrixAdapter):
        self._adapter = adapter

    def normalize(self, event: Any, room_id: str) -> ChannelMessage | None:
        if event.sender == self._adapter.user_id:
            return None

        content_type, text_content, attachments = self._extract_content(event)
        msg_type = self._map_message_type(content_type)
        chat_type = self._adapter._session.resolve_chat_type(room_id)
        is_thread, root_event_id = is_thread_event(event)

        mentions = self._extract_mentions(event)
        is_edit = self._is_edit_event(event)

        if chat_type == "group" and not is_thread:
            require_mention = get_room_require_mention(room_id, self._adapter.config)
            if require_mention and not (mentions and mentions.is_bot_mentioned):
                return None

        server_ts = getattr(event, "server_timestamp", None)
        timestamp = datetime.fromtimestamp((server_ts / 1000) if server_ts else datetime.now(UTC).timestamp(), tz=UTC)

        event_type = EventType.MESSAGE_UPDATED if is_edit else EventType.MESSAGE_RECEIVED

        return ChannelMessage(
            identity=ChannelIdentity(
                channel_id="matrix",
                channel_type=ChannelType.MATRIX,
                channel_user_id=event.sender,
                channel_chat_id=self._adapter._session.resolve_chat_id(room_id, chat_type == "direct"),
                channel_message_id=event.event_id,
            ),
            event_type=event_type,
            message_type=msg_type,
            chat_type=ChatType.THREAD if is_thread else (ChatType.DIRECT if chat_type == "direct" else ChatType.GROUP),
            content=text_content,
            attachments=attachments,
            mentions=mentions,
            reply_to_message_id=self._extract_reply_to(event),
            metadata={
                "room_id": room_id,
                "is_thread": is_thread,
                "root_event_id": root_event_id,
                "origin_server_ts": server_ts,
                "is_edit": is_edit,
            },
            timestamp=timestamp,
        )

    def _extract_content(self, event: Any) -> tuple[str, str, list[Attachment]]:
        if isinstance(event, RoomMessageFormatted):
            return "text", event.body or event.formatted_body or "", []
        if isinstance(event, RoomMessageText):
            return "text", event.body, []

        if isinstance(event, RoomMessageImage):
            img = event.source.get("content", {})
            mxc_url = img.get("url", "")
            return (
                "image",
                img.get("body", ""),
                [
                    Attachment(
                        type="image",
                        url=mxc_url,
                        mime_type=img.get("info", {}).get("mimetype"),
                    )
                ],
            )

        if isinstance(event, RoomMessageFile):
            file_content = event.source.get("content", {})
            mxc_url = file_content.get("url", "")
            return (
                "file",
                file_content.get("body", ""),
                [
                    Attachment(
                        type="file",
                        url=mxc_url,
                        filename=file_content.get("body"),
                        mime_type=file_content.get("info", {}).get("mimetype"),
                    )
                ],
            )

        if isinstance(event, RoomMessageAudio):
            media_content = event.source.get("content", {})
            mxc_url = media_content.get("url", "")
            return (
                "audio",
                media_content.get("body", ""),
                [
                    Attachment(
                        type="audio",
                        url=mxc_url,
                        mime_type=media_content.get("info", {}).get("mimetype"),
                    )
                ],
            )

        if isinstance(event, RoomMessageVideo):
            media_content = event.source.get("content", {})
            mxc_url = media_content.get("url", "")
            return (
                "video",
                media_content.get("body", ""),
                [
                    Attachment(
                        type="video",
                        url=mxc_url,
                        mime_type=media_content.get("info", {}).get("mimetype"),
                    )
                ],
            )

        event_source = getattr(event, "source", {})
        location = parse_location_event(event_source)
        if location:
            return "location", location["text"], [Attachment(type="location", url=location.get("geo_uri", ""))]

        if is_sticker_event(event_source):
            sticker = parse_sticker(event_source)
            if sticker:
                return (
                    "sticker",
                    sticker["text"],
                    [Attachment(type="sticker", url=sticker.get("url", ""), mime_type=sticker.get("mimetype", ""))],
                )

        return "text", "", []

    def _map_message_type(self, content_type: str) -> MessageType:
        mapping = {
            "text": MessageType.TEXT,
            "image": MessageType.IMAGE,
            "video": MessageType.VIDEO,
            "audio": MessageType.AUDIO,
            "file": MessageType.FILE,
            "location": MessageType.TEXT,
            "sticker": MessageType.IMAGE,
        }
        return mapping.get(content_type, MessageType.TEXT)

    def _extract_mentions(self, event: Any) -> MentionsInfo | None:
        content = event.source.get("content", {})
        formatted_body = content.get("formatted_body", "")
        body = content.get("body", "")

        if not formatted_body and not body:
            return None

        is_at_bot = self._adapter.user_id in formatted_body or self._adapter.user_id in body

        mentioned = _MXID_PATTERN.findall(formatted_body) + _MXID_URL_PATTERN.findall(formatted_body)

        if not mentioned:
            plain_body = body if body else formatted_body
            plain_mentions = _PLAIN_MXID_PATTERN.findall(plain_body)
            mentioned = [f"@{local}:{domain}" for local, domain in plain_mentions]

        mentioned = list(dict.fromkeys(mentioned))

        return MentionsInfo(
            mentioned_user_ids=mentioned,
            is_bot_mentioned=is_at_bot,
        )

    def _extract_reply_to(self, event: Any) -> str | None:
        content = event.source.get("content", {})
        relates_to = content.get("m.relates_to", {})
        if relates_to.get("m.in_reply_to"):
            return relates_to["m.in_reply_to"].get("event_id")
        return None

    def _is_edit_event(self, event: Any) -> bool:
        content = event.source.get("content", {})
        relates_to = content.get("m.relates_to", {})
        return relates_to.get("rel_type") == "m.replace" and "m.new_content" in content