Support rendering message edits and deletions

4 files changed, 183 insertions, 18 deletions

diff --git a/slack/config.py b/slack/config.py
index 7527190..0b717fa 100644
--- a/slack/config.py
+++ b/slack/config.py
@@ -29,6 +29,13 @@ class SlackConfigSectionColor:
             WeeChatColor("blue"),
         )
 
+        self.deleted_message = WeeChatOption(
+            self._section,
+            "deleted_message",
+            "text color for a deleted message",
+            WeeChatColor("red"),
+        )
+
         self.disconnected = WeeChatOption(
             self._section,
             "disconnected",
@@ -36,6 +43,13 @@ class SlackConfigSectionColor:
             WeeChatColor("red"),
         )
 
+        self.edited_message_suffix = WeeChatOption(
+            self._section,
+            "edited_message_suffix",
+            "text color for the suffix after an edited message",
+            WeeChatColor("095"),
+        )
+
         self.loading = WeeChatOption(
             self._section,
             "loading",
diff --git a/slack/slack_conversation.py b/slack/slack_conversation.py
index 3b06a4d..f93dc7b 100644
--- a/slack/slack_conversation.py
+++ b/slack/slack_conversation.py
@@ -14,6 +14,7 @@ from slack.util import get_callback_name
 
 if TYPE_CHECKING:
     from slack_api.slack_conversations_info import SlackConversationsInfo
+    from slack_rtm.slack_rtm_message import SlackMessageChanged, SlackMessageDeleted
     from typing_extensions import Literal
 
     from slack.slack_api import SlackApi
@@ -36,6 +37,109 @@ def invalidate_nicklists():
             conversation.nicklist_needs_refresh = True
 
 
+def hdata_line_ts(line_pointer: str) -> Optional[SlackTs]:
+    data = weechat.hdata_pointer(weechat.hdata_get("line"), line_pointer, "data")
+    for i in range(
+        weechat.hdata_integer(weechat.hdata_get("line_data"), data, "tags_count")
+    ):
+        tag = weechat.hdata_string(
+            weechat.hdata_get("line_data"), data, f"{i}|tags_array"
+        )
+        if tag.startswith("slack_ts_"):
+            return SlackTs(tag[9:])
+    return None
+
+
+def tags_set_notify_none(tags: List[str]) -> List[str]:
+    notify_tags = {"notify_highlight", "notify_message", "notify_private"}
+    tags = [tag for tag in tags if tag not in notify_tags]
+    tags += ["no_highlight", "notify_none"]
+    return tags
+
+
+def modify_buffer_line(buffer_pointer: str, ts: SlackTs, new_text: str):
+    own_lines = weechat.hdata_pointer(
+        weechat.hdata_get("buffer"), buffer_pointer, "own_lines"
+    )
+    line_pointer = weechat.hdata_pointer(
+        weechat.hdata_get("lines"), own_lines, "last_line"
+    )
+
+    # Find the last line with this ts
+    is_last_line = True
+    while line_pointer and hdata_line_ts(line_pointer) != ts:
+        is_last_line = False
+        line_pointer = weechat.hdata_move(weechat.hdata_get("line"), line_pointer, -1)
+
+    if not line_pointer:
+        return
+
+    if shared.weechat_version >= 0x04000000:
+        data = weechat.hdata_pointer(weechat.hdata_get("line"), line_pointer, "data")
+        weechat.hdata_update(
+            weechat.hdata_get("line_data"), data, {"message": new_text}
+        )
+        return
+
+    # Find all lines for the message
+    pointers: List[str] = []
+    while line_pointer and hdata_line_ts(line_pointer) == ts:
+        pointers.append(line_pointer)
+        line_pointer = weechat.hdata_move(weechat.hdata_get("line"), line_pointer, -1)
+    pointers.reverse()
+
+    if not pointers:
+        return
+
+    if is_last_line:
+        lines = new_text.split("\n")
+        extra_lines_count = len(lines) - len(pointers)
+        if extra_lines_count > 0:
+            line_data = weechat.hdata_pointer(
+                weechat.hdata_get("line"), pointers[0], "data"
+            )
+            tags_count = weechat.hdata_integer(
+                weechat.hdata_get("line_data"), line_data, "tags_count"
+            )
+            tags = [
+                weechat.hdata_string(
+                    weechat.hdata_get("line_data"), line_data, f"{i}|tags_array"
+                )
+                for i in range(tags_count)
+            ]
+            tags = tags_set_notify_none(tags)
+            tags_str = ",".join(tags)
+            last_read_line = weechat.hdata_pointer(
+                weechat.hdata_get("lines"), own_lines, "last_read_line"
+            )
+            should_set_unread = last_read_line == pointers[-1]
+
+            # Insert new lines to match the number of lines in the message
+            weechat.buffer_set(buffer_pointer, "print_hooks_enabled", "0")
+            for _ in range(extra_lines_count):
+                weechat.prnt_date_tags(buffer_pointer, ts.major, tags_str, " \t ")
+                pointers.append(
+                    weechat.hdata_pointer(
+                        weechat.hdata_get("lines"), own_lines, "last_line"
+                    )
+                )
+            if should_set_unread:
+                weechat.buffer_set(buffer_pointer, "unread", "")
+            weechat.buffer_set(buffer_pointer, "print_hooks_enabled", "1")
+    else:
+        # Split the message into at most the number of existing lines as we can't insert new lines
+        lines = new_text.split("\n", len(pointers) - 1)
+        # Replace newlines to prevent garbled lines in bare display mode
+        lines = [line.replace("\n", " | ") for line in lines]
+
+    # Extend lines in case the new message is shorter than the old as we can't delete lines
+    lines += [""] * (len(pointers) - len(lines))
+
+    for pointer, line in zip(pointers, lines):
+        data = weechat.hdata_pointer(weechat.hdata_get("line"), pointer, "data")
+        weechat.hdata_update(weechat.hdata_get("line_data"), data, {"message": line})
+
+
 class SlackConversation:
     def __init__(
         self,
@@ -268,6 +372,20 @@ class SlackConversation:
                 f"{self.buffer_pointer};off;{user.nick()}",
             )
 
+    async def change_message(self, data: SlackMessageChanged):
+        ts = SlackTs(data["ts"])
+        message = self._messages.get(ts)
+        if message:
+            message.update_message_json(data["message"])
+            modify_buffer_line(self.buffer_pointer, ts, await message.render_message())
+
+    async def delete_message(self, data: SlackMessageDeleted):
+        ts = SlackTs(data["deleted_ts"])
+        message = self._messages.get(ts)
+        if message:
+            message.deleted = True
+            modify_buffer_line(self.buffer_pointer, ts, await message.render_message())
+
     async def typing_add_user(self, user_id: str, thread_ts: Optional[str]):
         if not shared.config.look.typing_status_nicks.value:
             return
diff --git a/slack/slack_message.py b/slack/slack_message.py
index 403dc41..6fb6c37 100644
--- a/slack/slack_message.py
+++ b/slack/slack_message.py
@@ -46,9 +46,11 @@ class SlackTs(str):
 class SlackMessage:
     def __init__(self, conversation: SlackConversation, message_json: SlackMessageDict):
         self._message_json = message_json
-        self._rendered = None
+        self._rendered_prefix = None
+        self._rendered_message = None
         self.conversation = conversation
         self.ts = SlackTs(message_json["ts"])
+        self._deleted = False
 
     @property
     def workspace(self) -> SlackWorkspace:
@@ -75,6 +77,20 @@ class SlackMessage:
     def priority(self) -> MessagePriority:
         return MessagePriority.MESSAGE
 
+    @property
+    def deleted(self) -> bool:
+        return self._deleted
+
+    @deleted.setter
+    def deleted(self, value: bool):
+        self._deleted = value
+        self._rendered_message = None
+
+    def update_message_json(self, message_json: SlackMessageDict):
+        self._message_json = message_json
+        self._rendered_prefix = None
+        self._rendered_message = None
+
     async def tags(self, backlog: bool = False) -> str:
         nick = await self._nick(colorize=False, only_nick=True)
         tags = [f"slack_ts_{self.ts}", f"nick_{nick}"]
@@ -118,11 +134,8 @@ class SlackMessage:
         return ",".join(tags)
 
     async def render(self) -> str:
-        if self._rendered is not None:
-            return self._rendered
-
-        prefix_coro = self._prefix()
-        message_coro = self._render_message()
+        prefix_coro = self.render_prefix()
+        message_coro = self.render_message()
         prefix, message = await gather(prefix_coro, message_coro)
         self._rendered = f"{prefix}\t{message}"
         return self._rendered
@@ -142,7 +155,9 @@ class SlackMessage:
             user = await self.workspace.users[self._message_json["user"]]
             return user.nick(colorize=colorize, only_nick=only_nick)
 
-    async def _prefix(self, colorize: bool = True, only_nick: bool = False) -> str:
+    async def _render_prefix(
+        self, colorize: bool = True, only_nick: bool = False
+    ) -> str:
         if self._message_json.get("subtype") in ["channel_join", "group_join"]:
             return removesuffix(weechat.prefix("join"), "\t")
         elif self._message_json.get("subtype") in ["channel_leave", "group_leave"]:
@@ -150,8 +165,16 @@ class SlackMessage:
         else:
             return await self._nick(colorize=colorize, only_nick=only_nick)
 
+    async def render_prefix(self) -> str:
+        if self._rendered_prefix is not None:
+            return self._rendered_prefix
+        self._rendered_prefix = await self._render_prefix()
+        return self._rendered_prefix
+
     async def _render_message(self) -> str:
-        if self._message_json.get("subtype") in [
+        if self._deleted:
+            return with_color(shared.config.color.deleted_message.value, "(deleted)")
+        elif self._message_json.get("subtype") in [
             "channel_join",
             "group_join",
             "channel_leave",
@@ -180,7 +203,19 @@ class SlackMessage:
 
             return f"{await self._nick()} {text_action} {text_conversation_name}{inviter_text}"
         else:
-            return await self._unfurl_refs(self._message_json["text"])
+            text = await self._unfurl_refs(self._message_json["text"])
+            text_edited = (
+                f" {with_color(shared.config.color.edited_message_suffix.value, '(edited)')}"
+                if self._message_json.get("edited")
+                else ""
+            )
+            return text + text_edited
+
+    async def render_message(self) -> str:
+        if self._rendered_message is not None:
+            return self._rendered_message
+        self._rendered_message = await self._render_message()
+        return self._rendered_message
 
     def _item_prefix(self, item_id: str):
         if item_id.startswith("#") or item_id.startswith("@"):
diff --git a/slack/slack_workspace.py b/slack/slack_workspace.py
index 0d9405f..4b045e1 100644
--- a/slack/slack_workspace.py
+++ b/slack/slack_workspace.py
@@ -279,20 +279,18 @@ class SlackWorkspace:
             else:
                 channel = None
 
-            if data["type"] == "message":
+            if data["type"] == "message" and channel is not None:
                 if "subtype" in data and data["subtype"] == "message_changed":
-                    pass
+                    await channel.change_message(data)
                 elif "subtype" in data and data["subtype"] == "message_deleted":
-                    pass
+                    await channel.delete_message(data)
                 elif "subtype" in data and data["subtype"] == "message_replied":
                     pass
                 else:
-                    if channel:
-                        message = SlackMessage(channel, data)
-                        await channel.add_message(message)
-            elif data["type"] == "user_typing":
-                if channel:
-                    await channel.typing_add_user(data["user"], data.get("thread_ts"))
+                    message = SlackMessage(channel, data)
+                    await channel.add_message(message)
+            elif data["type"] == "user_typing" and channel is not None:
+                await channel.typing_add_user(data["user"], data.get("thread_ts"))
             else:
                 weechat.prnt("", f"\t{self.name} received: {json.dumps(data)}")
         except Exception as e: