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

import asyncio
import base64
import json
import logging
from datetime import datetime, UTC
from typing import Any

from yuxi.channel.context import ChannelContext
from yuxi.channel.extensions.signal.access_policy import (
    check_dm_access,
    check_group_access,
    generate_pairing_code,
    resolve_group_config,
)
from yuxi.channel.extensions.signal.client import SignalRpcClient, SignalSseClient
from yuxi.channel.extensions.signal.dedupe import (
    SignalDedupeStore,
    build_signal_dedupe_key,
)
from yuxi.channel.extensions.signal.identity import (
    SignalSender,
    SignalSenderPhone,
    format_signal_sender_id,
    resolve_signal_peer_id,
    resolve_signal_recipient,
    resolve_signal_sender,
)
from yuxi.channel.extensions.signal.send import send_signal_receipt
from yuxi.channel.extensions.signal.sse_reconnect import SseReconnector
from yuxi.channel.message.media_store import MediaStore
from yuxi.channel.message.models import (
    GroupContext,
    MessageType,
    PeerInfo,
    UnifiedMessage,
)
from yuxi.channel.routing.models import PeerKind

logger = logging.getLogger(__name__)


class SignalEventProcessor:
    def __init__(
        self,
        rpc_client: SignalRpcClient | None,
        sse_client: SignalSseClient | None,
        account_config: dict,
        dedupe_store: SignalDedupeStore,
        reconnector: SseReconnector,
        media_store: MediaStore | None = None,
    ):
        self._rpc_client = rpc_client
        self._sse_client = sse_client
        self._account_config = account_config
        self._dedupe_store = dedupe_store
        self._reconnector = reconnector
        self._media_store = media_store

        self._message_registry: dict[str, dict] = {}
        self._group_history: dict[str, list[dict]] = {}
        self._pairing_store: dict[str, str] = {}
        self._msg_id_counter = 0

        self._receipt_event_index = 0
        self._typing_event_index = 0
        self._call_event_index = 0
        self._remote_delete_event_index = 0
        self._sync_event_index = 0
        self._reaction_event_index = 0

    @property
    def message_registry(self) -> dict[str, dict]:
        return self._message_registry

    @property
    def pairing_store(self) -> dict[str, str]:
        return self._pairing_store

    @property
    def msg_id_counter(self) -> int:
        return self._msg_id_counter

    async def run_event_loop(self, ctx: ChannelContext) -> None:
        if not self._sse_client:
            return

        account_phone = self._account_config.get("account", "")

        while not ctx.cancel_event.is_set():
            try:
                async for sse_event in self._sse_client.stream_events(ctx.cancel_event):
                    self._reconnector.reset()
                    await self._process_sse_event(sse_event, ctx, account_phone)
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("Signal SSE stream error")
                await self._reconnector.wait_before_reconnect(ctx.cancel_event)

    async def _process_sse_event(self, sse_event, ctx: ChannelContext, account_phone: str):
        try:
            payload = json.loads(sse_event.data)
        except json.JSONDecodeError:
            return

        envelope = payload.get("envelope", payload)
        source_number = envelope.get("sourceNumber") or envelope.get("source")
        source_uuid = envelope.get("sourceUuid")
        source_name = envelope.get("sourceName")

        sender = resolve_signal_sender(source_number, source_uuid)
        if not sender:
            return

        sender_id = format_signal_sender_id(sender)
        account_uuid = self._account_config.get("account_uuid", "")

        if source_uuid and account_uuid and source_uuid == account_uuid:
            return
        if isinstance(sender, SignalSenderPhone) and sender.raw == account_phone:
            return

        receipt_message = envelope.get("receiptMessage")
        if receipt_message:
            await self._handle_receipt_message(ctx, receipt_message, sender_id)
            return

        typing_message = envelope.get("typingMessage")
        if typing_message:
            await self._handle_typing_message(ctx, typing_message, sender_id)
            return

        call_message = envelope.get("callMessage")
        if call_message:
            await self._handle_call_message(ctx, call_message, sender_id)
            return

        sync_message = envelope.get("syncMessage")
        if sync_message:
            await self._handle_sync_message(ctx, sync_message, sender_id)
            return

        data_message = envelope.get("dataMessage")
        edit_message = envelope.get("editMessage")
        dm = edit_message.get("dataMessage") if edit_message else data_message

        if not dm:
            return

        text = dm.get("message", "") or ""
        timestamp_val = dm.get("timestamp") or envelope.get("timestamp", 0)

        reaction = dm.get("reaction")
        if reaction and not text.strip():
            await self._handle_reaction_notification(ctx, reaction, sender_id, envelope)
            return

        remote_delete = dm.get("remoteDelete")
        if remote_delete:
            await self._handle_remote_delete(ctx, remote_delete, sender_id)
            return

        group_info = envelope.get("groupInfo") or dm.get("groupInfo")
        group_id = group_info.get("groupId", "") if group_info else ""
        is_group = bool(group_id)

        if is_group:
            group_cfg = resolve_group_config(self._account_config, group_id)
            access_status, reason = check_group_access(
                self._account_config,
                sender_id,
                group_allow_from=self._account_config.get("group_allow_from"),
                account_key=ctx.account_id,
            )
            if access_status == "deny":
                if group_cfg.get("ingest"):
                    self._append_group_history(group_id, sender_id, text)
                return

            if group_cfg.get("require_mention", True):
                mentions = dm.get("mentions", [])
                if not self._is_mentioned(mentions, account_phone, account_uuid):
                    if group_cfg.get("ingest"):
                        self._append_group_history(group_id, sender_id, text)
                    return
        else:
            access_status, reason = check_dm_access(
                self._account_config,
                sender_id,
                pairing_store=self._pairing_store,
            )
            if access_status == "deny":
                return
            if access_status == "pairing":
                await self._handle_pairing_request(ctx, sender, sender_id)
                return

        self._msg_id_counter += 1

        dedupe_key = build_signal_dedupe_key(ctx.account_id, timestamp_val, sender_id)
        if self._dedupe_store.is_duplicate(dedupe_key):
            return
        self._dedupe_store.mark_seen(dedupe_key)

        peer_kind = PeerKind.GROUP if is_group else PeerKind.DIRECT
        peer_info = PeerInfo(
            kind=peer_kind,
            id=resolve_signal_peer_id(sender),
            display_name=source_name or sender_id,
            username=sender_id,
        )

        group_ctx = None
        if is_group and group_id:
            group_ctx = GroupContext(
                id=group_id,
                name=group_info.get("name") if group_info else None,
            )

        reply_to_id = None
        quote = dm.get("quote")
        if quote:
            reply_to_id = str(quote.get("id", ""))

        attachments = dm.get("attachments", [])
        media_urls: list[str] = []
        image_base64: str | None = None
        media_types: list[str] = []

        if attachments:
            for att in attachments:
                att_id = att.get("id", "")
                content_type = att.get("contentType", "")
                filename = att.get("filename", "")
                if not att_id:
                    continue

                media_types.append(content_type)

                if self._media_store and self._rpc_client:
                    try:
                        raw = await self._rpc_client.get_attachment(
                            att_id,
                            account=self._account_config.get("account"),
                            recipient=resolve_signal_recipient(sender),
                        )

                        if content_type.startswith("image/") and len(raw) <= 5 * 1_048_576:
                            image_base64 = base64.b64encode(raw).decode()
                            media_urls.append(f"signal-attachment:{att_id}")
                        else:
                            stored_id = await self._media_store.store(
                                data=raw,
                                content_type=content_type,
                                filename=filename,
                                source_url=f"signal-attachment:{att_id}",
                            )
                            if stored_id:
                                media_urls.append(f"media-store:{stored_id}")
                            else:
                                media_urls.append(f"signal-attachment:{att_id}")
                    except Exception:
                        logger.debug("Failed to download Signal attachment %s", att_id)
                        media_urls.append(f"signal-attachment:{att_id}")
                else:
                    media_urls.append(f"signal-attachment:{att_id}")

        try:
            ts = datetime.fromtimestamp(timestamp_val / 1000, tz=UTC) if timestamp_val else datetime.now(UTC)
        except Exception:
            ts = datetime.now(UTC)

        msg_id = str(self._msg_id_counter)

        self._message_registry[msg_id] = {
            "recipient": resolve_signal_recipient(sender),
            "timestamp": timestamp_val,
            "group_id": group_id if is_group else None,
            "sender_id": sender_id,
        }

        message_type = MessageType.TEXT
        if media_urls:
            if any(t.startswith("image/") for t in media_types):
                message_type = MessageType.IMAGE
            elif any(t.startswith("audio/") for t in media_types):
                message_type = MessageType.VOICE
            else:
                message_type = MessageType.FILE

        metadata: dict[str, Any] = {}
        expires_in = dm.get("expiresInSeconds")
        if expires_in:
            metadata["expires_in_seconds"] = expires_in
        if dm.get("isViewOnce"):
            metadata["is_view_once"] = True
        sticker = dm.get("sticker")
        if sticker:
            metadata["sticker"] = sticker
        preview = dm.get("preview")
        if preview:
            metadata["preview"] = preview
        contact = dm.get("contact")
        if contact:
            metadata["contact"] = contact
        if group_info and group_info.get("type"):
            metadata["group_info_type"] = group_info.get("type")

        msg = UnifiedMessage(
            msg_id=msg_id,
            channel_type="signal",
            account_id=ctx.account_id,
            content=text,
            sender=peer_info,
            message_type=message_type,
            media_urls=media_urls,
            image_base64=image_base64,
            media_types=media_types,
            group=group_ctx,
            timestamp=ts,
            raw_payload=payload,
            reply_to_id=reply_to_id,
            metadata=metadata,
        )

        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                logger.exception("Failed to enqueue Signal message")

        if not is_group and self._rpc_client:
            auto_start = self._account_config.get("auto_start", True)
            daemon_level_receipts = auto_start and self._account_config.get("send_read_receipts", False)
            program_level_receipts = self._account_config.get("send_read_receipts", True) and not daemon_level_receipts
            if program_level_receipts:
                try:
                    await send_signal_receipt(
                        self._rpc_client,
                        resolve_signal_recipient(sender),
                        timestamp_val,
                        account=self._account_config.get("account"),
                    )
                except Exception:
                    logger.debug("Failed to send read receipt for DM from %s", sender_id)

    def _is_mentioned(self, mentions, account_phone, account_uuid):
        for m in mentions:
            if m.get("number") == account_phone or m.get("uuid") == account_uuid:
                return True
        return False

    async def _handle_receipt_message(self, ctx: ChannelContext, receipt_message, sender_id):
        reaction_notifications = self._account_config.get("reaction_notifications", "own")
        if reaction_notifications == "off":
            return

        receipts = receipt_message.get("timestamps", [])
        receipt_type = receipt_message.get("type", "unknown")
        if not receipts:
            return

        content = f"receipt {receipt_type}: message timestamps {receipts} from {sender_id}"

        self._receipt_event_index += 1

        msg = UnifiedMessage(
            msg_id=f"receipt-{self._receipt_event_index}",
            channel_type="signal",
            account_id=ctx.account_id,
            content=content,
            sender=PeerInfo(kind=PeerKind.DIRECT, id=sender_id, display_name=sender_id),
            message_type=MessageType.EVENT,
            timestamp=datetime.now(UTC),
            raw_payload={"receiptMessage": receipt_message},
        )
        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                pass

    async def _handle_typing_message(self, ctx: ChannelContext, typing_message, sender_id):
        if not self._account_config.get("typing_notifications", False):
            return

        action = "stopped" if typing_message.get("action") == "STOPPED" else "started"
        content = f"typing {action} by {sender_id}"

        self._typing_event_index += 1

        msg = UnifiedMessage(
            msg_id=f"typing-{self._typing_event_index}",
            channel_type="signal",
            account_id=ctx.account_id,
            content=content,
            sender=PeerInfo(kind=PeerKind.DIRECT, id=sender_id, display_name=sender_id),
            message_type=MessageType.EVENT,
            timestamp=datetime.now(UTC),
            raw_payload={"typingMessage": typing_message},
        )
        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                pass

    async def _handle_call_message(self, ctx: ChannelContext, call_message, sender_id):
        self._call_event_index += 1

        content = f"call event from {sender_id}"
        msg = UnifiedMessage(
            msg_id=f"call-{self._call_event_index}",
            channel_type="signal",
            account_id=ctx.account_id,
            content=content,
            sender=PeerInfo(kind=PeerKind.DIRECT, id=sender_id, display_name=sender_id),
            message_type=MessageType.EVENT,
            timestamp=datetime.now(UTC),
            raw_payload={"callMessage": call_message},
        )
        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                pass

    async def _handle_remote_delete(self, ctx: ChannelContext, remote_delete, sender_id):
        self._remote_delete_event_index += 1

        timestamp = remote_delete.get("timestamp", 0)
        content = f"remote delete: message timestamp {timestamp} by {sender_id}"
        msg = UnifiedMessage(
            msg_id=f"remote-delete-{self._remote_delete_event_index}",
            channel_type="signal",
            account_id=ctx.account_id,
            content=content,
            sender=PeerInfo(kind=PeerKind.DIRECT, id=sender_id, display_name=sender_id),
            message_type=MessageType.EVENT,
            timestamp=datetime.now(UTC),
            raw_payload={"remoteDelete": remote_delete},
        )
        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                pass

    async def _handle_sync_message(self, ctx: ChannelContext, sync_message, sender_id):
        self._sync_event_index += 1

        content = f"sync message from {sender_id}"
        msg = UnifiedMessage(
            msg_id=f"sync-{self._sync_event_index}",
            channel_type="signal",
            account_id=ctx.account_id,
            content=content,
            sender=PeerInfo(kind=PeerKind.DIRECT, id=sender_id, display_name=sender_id),
            message_type=MessageType.EVENT,
            timestamp=datetime.now(UTC),
            raw_payload={"syncMessage": sync_message},
        )
        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                pass

    def _append_group_history(self, group_id: str, sender_id: str, text: str) -> None:
        history_limit = self._account_config.get("history_limit", 50)
        buf = self._group_history.setdefault(group_id, [])
        buf.append({"sender": sender_id, "text": text})
        if len(buf) > history_limit:
            buf.pop(0)

    async def _handle_pairing_request(self, ctx: ChannelContext, sender: SignalSender, sender_id: str):
        code = generate_pairing_code()
        self._pairing_store[sender_id] = code

        recipient = resolve_signal_recipient(sender)
        if self._rpc_client:
            from yuxi.channel.extensions.signal.send import send_signal_message

            try:
                await send_signal_message(
                    self._rpc_client,
                    recipient,
                    f"To start chatting, please reply with this pairing code: {code}",
                    account=self._account_config.get("account"),
                )
            except Exception:
                logger.exception("Failed to send pairing challenge")

        logger.info("Pairing challenge sent to %s", sender_id)

    async def _handle_reaction_notification(self, ctx: ChannelContext, reaction, sender_id, envelope):
        reaction_notifications = self._account_config.get("reaction_notifications", "own")
        if reaction_notifications == "off":
            return

        emoji = reaction.get("emoji", "")
        is_remove = reaction.get("remove", False)
        target_author = reaction.get("targetAuthorNumber") or reaction.get("targetAuthorUuid", "")

        action = "removed" if is_remove else "added"
        content = f"reaction {action}: {emoji} by {sender_id}"
        if target_author:
            content += f" on message from {target_author}"

        self._reaction_event_index += 1

        msg = UnifiedMessage(
            msg_id=f"reaction-{self._reaction_event_index}",
            channel_type="signal",
            account_id=ctx.account_id,
            content=content,
            sender=PeerInfo(kind=PeerKind.DIRECT, id=sender_id, display_name=sender_id),
            message_type=MessageType.EVENT,
            timestamp=datetime.now(UTC),
            raw_payload={"reaction": reaction, "envelope": envelope},
        )

        if ctx.queue:
            try:
                await ctx.queue.put(msg)
            except Exception:
                pass