From fb9060555775ef2237ef9eb045274facb5ace567 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Thu, 8 Feb 2018 14:22:48 +0200
Subject: [PATCH] Fix replying to real Matrix users messages in a multiuser
 setup

---
 mautrix_telegram/formatter.py | 12 ++++++------
 mautrix_telegram/portal.py    | 31 ++++++++++++++++++++++++-------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/mautrix_telegram/formatter.py b/mautrix_telegram/formatter.py
index f1a69375..030f72aa 100644
--- a/mautrix_telegram/formatter.py
+++ b/mautrix_telegram/formatter.py
@@ -95,12 +95,12 @@ class MatrixParser(HTMLParser):
                     entity_type = MessageEntityMentionName
                     args["user_id"] = user.tgid
             elif reply and self._tg_space and (len(self.entities) == 0
-                                              and len(self._building_entities) == 0):
+                                               and len(self._building_entities) == 0):
                 room_id = reply.group(1)
                 message_id = reply.group(2)
-                message = DBMessage.query.filter(DBMessage.mxid == message_id
-                                                 and DBMessage.mx_room == room_id
-                                                 and DBMessage.tg_space == self._tg_space
+                message = DBMessage.query.filter(DBMessage.mxid == message_id,
+                                                 DBMessage.mx_room == room_id,
+                                                 DBMessage.tg_space == self._tg_space
                                                  ).one_or_none()
                 if not message:
                     return
@@ -221,7 +221,8 @@ def telegram_event_to_matrix(evt, source):
                 + f"<blockquote>{html}</blockquote>")
 
     if evt.reply_to_msg_id:
-        space = evt.to_id.channel_id if isinstance(evt, Message) and isinstance(evt.to_id, PeerChannel) else source.tgid
+        space = evt.to_id.channel_id if isinstance(evt, Message) and isinstance(evt.to_id,
+                                                                                PeerChannel) else source.tgid
         msg = DBMessage.query.get((evt.reply_to_msg_id, space))
         if msg:
             quote = f"<a href=\"https://matrix.to/#/{msg.mx_room}/{msg.mxid}\">Quote<br></a>"
@@ -315,5 +316,4 @@ def _telegram_to_matrix(text, entities):
 
     return "".join(html)
 
-
 # endregion
diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py
index 2f7f754e..008ea49a 100644
--- a/mautrix_telegram/portal.py
+++ b/mautrix_telegram/portal.py
@@ -17,6 +17,7 @@
 from io import BytesIO
 from collections import deque
 from datetime import datetime
+import random
 import mimetypes
 import hashlib
 import logging
@@ -57,6 +58,7 @@ class Portal:
         self._main_intent = None
 
         self._dedup = deque()
+        self._dedup_mxid = {}
 
         if tgid:
             self.by_tgid[self.tgid_full] = self
@@ -110,13 +112,16 @@ class Portal:
                            .encode("utf-8")
                            ).hexdigest()
 
-    def is_duplicate(self, event):
+    def is_duplicate(self, event, mxid=None):
         hash = self._hash_event(event)
         if hash in self._dedup:
-            return True
+            return self._dedup_mxid[hash]
+
+        self._dedup_mxid[hash] = mxid
         self._dedup.append(hash)
+
         if len(self._dedup) > 20:
-            self._dedup.popleft()
+            del self._dedup_mxid[self._dedup.popleft()]
         return False
 
     def get_input_entity(self, user):
@@ -381,6 +386,7 @@ class Portal:
         if type in {"m.text", "m.emote"}:
             if "format" in message and message["format"] == "org.matrix.custom.html":
                 space = self.tgid if self.peer_type == "channel" else sender.tgid
+                print(sender.username, sender.tgid, space)
                 message, entities = formatter.matrix_to_telegram(message["formatted_body"], space)
                 if type == "m.emote":
                     message = "/me " + message
@@ -410,7 +416,7 @@ class Portal:
         else:
             self.log.debug("Unhandled Matrix event: %s", message)
             return
-        self.is_duplicate(response)
+        self.is_duplicate(response, event_id)
         self.db.add(DBMessage(
             tgid=response.id,
             tg_space=self.tgid if self.peer_type == "channel" else sender.tgid,
@@ -681,7 +687,13 @@ class Portal:
         if not self.mxid:
             self.create_matrix_room(source, invites=[source.mxid])
 
-        if self.is_duplicate(evt):
+        tg_space = self.tgid if self.peer_type == "channel" else source.tgid
+
+        temporary_identifier = f"${random.randint(1000000000000,9999999999999)}TGBRIDGETEMP"
+        mxid = self.is_duplicate(evt, temporary_identifier)
+        if mxid:
+            self.db.add(DBMessage(tgid=evt.id, mx_room=self.mxid, mxid=mxid, tg_space=tg_space))
+            self.db.commit()
             return
 
         if evt.message:
@@ -700,8 +712,13 @@ class Portal:
             self.log.debug("Unhandled Telegram message: %s", evt)
             return
 
-        self.db.add(DBMessage(tgid=evt.id, mx_room=self.mxid, mxid=response["event_id"],
-                              tg_space=self.tgid if self.peer_type == "channel" else source.tgid))
+        mxid = response["event_id"]
+        DBMessage.query \
+            .filter(DBMessage.mx_room == self.mxid,
+                    DBMessage.mxid == temporary_identifier,
+                    DBMessage.tg_space == tg_space) \
+            .update({"mxid": mxid})
+        self.db.add(DBMessage(tgid=evt.id, mx_room=self.mxid, mxid=mxid, tg_space=tg_space))
         self.db.commit()
 
     def handle_telegram_action(self, source, sender, action):