From a215b0bc65d26cffb07e6a60a0efcc1c4ab0ab74 Mon Sep 17 00:00:00 2001
From: Paul Nameless <reacsdas@gmail.com>
Date: Thu, 21 May 2020 21:55:57 +0800
Subject: [PATCH] Add forward msg functionality and multiple msgs selection

---
 tg/controllers.py | 41 +++++++++++++++++++++++++++++++++++++++++
 tg/models.py      |  2 ++
 tg/msg.py         |  6 ++++++
 tg/tdlib.py       | 24 ++++++++++++++++++++++++
 tg/views.py       |  7 +++++++
 5 files changed, 80 insertions(+)

diff --git a/tg/controllers.py b/tg/controllers.py
index 2dd1cbe..d283591 100644
--- a/tg/controllers.py
+++ b/tg/controllers.py
@@ -97,8 +97,49 @@ class Controller:
             "I": self.write_long_msg,
             "A": self.write_long_msg,
             "bp": self.breakpoint,
+
+            " ": self.select_msg,
+            "^[": self.discard_selected_msgs,  # esc
+            "y": self.copy_msgs,
+            "p": self.forward_msgs,
         }
 
+    def forward_msgs(self, _: int):
+        chat_id = self.model.chats.id_by_index(self.model.current_chat)
+        if not self.model.yanked_msgs:
+            return
+        from_chat_id, msg_ids = self.model.yanked_msgs
+        self.tg.forward_msgs(chat_id, from_chat_id, msg_ids)
+        self.present_info(f"Forwarded {len(msg_ids)} messages")
+
+    def copy_msgs(self, _: int):
+        # can_be_forwarded
+        chat_id = self.model.chats.id_by_index(self.model.current_chat)
+        msg_ids = self.model.selected[chat_id]
+        if not msg_ids:
+            self.present_error("No msgs selected")
+            return
+        self.model.yanked_msgs = (chat_id, msg_ids)
+        self.discard_selected_msgs(0)
+        self.present_info(f"Copied {len(msg_ids)} messages")
+
+    def select_msg(self, _: int):
+        chat_id = self.model.chats.id_by_index(self.model.current_chat)
+        msg = MsgProxy(self.model.current_msg)
+
+        if msg.msg_id in self.model.selected[chat_id]:
+            self.model.selected[chat_id].remove(msg.msg_id)
+        else:
+            self.model.selected[chat_id].append(msg.msg_id)
+        self.model.next_msg(1)
+        self.refresh_msgs()
+        self.present_info("Removed selections")
+
+    def discard_selected_msgs(self, _: int):
+        chat_id = self.model.chats.id_by_index(self.model.current_chat)
+        self.model.selected[chat_id] = []
+        self.refresh_msgs()
+
     def jump_bottom(self, _: int):
         log.info("jump_bottom:")
         if self.model.jump_bottom():
diff --git a/tg/models.py b/tg/models.py
index c926ffd..2a15a83 100644
--- a/tg/models.py
+++ b/tg/models.py
@@ -16,6 +16,8 @@ class Model:
         self.users = UserModel(tg)
         self.current_chat = 0
         self.downloads: Dict[int, Tuple[int, int]] = {}
+        self.selected: Dict[int, List[int]] = defaultdict(list)
+        self.yanked_msgs: Tuple[int, List[int]] = None
 
     def get_me(self):
         return self.users.get_me()
diff --git a/tg/msg.py b/tg/msg.py
index ba8474a..c0254c3 100644
--- a/tg/msg.py
+++ b/tg/msg.py
@@ -187,3 +187,9 @@ class MsgProxy:
     @property
     def sender_id(self) -> int:
         return self.msg["sender_user_id"]
+
+    @property
+    def forward(self) -> Optional[Dict[str, Any]]:
+        if "forward_info" not in self.msg:
+            return None
+        return self.msg["forward_info"]
diff --git a/tg/tdlib.py b/tg/tdlib.py
index 7fb789f..605efb0 100644
--- a/tg/tdlib.py
+++ b/tg/tdlib.py
@@ -1,3 +1,5 @@
+from typing import Any, Dict, List
+
 from telegram.client import AsyncResult, Telegram
 
 
@@ -149,3 +151,25 @@ class Tdlib(Telegram):
             "message_id": message_id,
         }
         return self._send_data(data)
+
+    def forward_msgs(
+        self,
+        chat_id: int,
+        from_chat_id: int,
+        message_ids: List[int],
+        as_album: bool = False,
+        send_copy: bool = False,
+        remove_caption: bool = False,
+        options: Dict[str, Any] = {},
+    ) -> AsyncResult:
+        data = {
+            "@type": "forwardMessages",
+            "chat_id": chat_id,
+            "from_chat_id": from_chat_id,
+            "message_ids": message_ids,
+            "as_album": as_album,
+            "send_copy": send_copy,
+            "remove_caption": remove_caption,
+            "options": options,
+        }
+        return self._send_data(data)
diff --git a/tg/views.py b/tg/views.py
index 5515146..f4cee95 100644
--- a/tg/views.py
+++ b/tg/views.py
@@ -259,6 +259,12 @@ class MsgView:
         flags = []
         chat = self.model.chats.chats[self.model.current_chat]
 
+        if msg_proxy.msg_id in self.model.selected[chat["id"]]:
+            flags.append("selected")
+
+        if msg_proxy.forward is not None:
+            flags.append("forwarded")
+
         if (
             not self.model.is_me(msg_proxy.sender_id)
             and msg_proxy.msg_id > chat["last_read_inbox_message_id"]
@@ -276,6 +282,7 @@ class MsgView:
             flags.append(self.states.get(state_type, state_type))
         if msg_proxy.msg["edit_date"]:
             flags.append("edited")
+
         if not flags:
             return ""
         return ",".join(flags)