From 9cadb9541a1f0e02f2739710d6e362abb2a19125 Mon Sep 17 00:00:00 2001
From: blank X <theblankx@protonmail.com>
Date: Fri, 4 Sep 2020 21:11:02 +0700
Subject: [PATCH] Allow for multi chat/subreddit support

---
 example_config.yaml | 12 ++++--
 redditbot.py        | 94 ++++++++++++++++++++++++++-------------------
 2 files changed, 64 insertions(+), 42 deletions(-)

diff --git a/example_config.yaml b/example_config.yaml
index 3daf324..dafac10 100644
--- a/example_config.yaml
+++ b/example_config.yaml
@@ -11,9 +11,15 @@ config:
     storage_chat: -1001222674489
     storage_message_id: 1367
     send_to_chats:
-        - goodanimemes
-    subreddits:
-        - goodanimemes
+        - goodanimemes:
+            - goodanimemes
+        - -1001327025649:
+            - ProgrammerAnimemes
+        - memes10k:
+            - linuxmemes
+            - dankmemes
+            - memes
+            - ShitPostCrusaders
     bot_admins:
         - thekneesocks
         - 214416808
diff --git a/redditbot.py b/redditbot.py
index 474c26b..4104780 100644
--- a/redditbot.py
+++ b/redditbot.py
@@ -31,8 +31,15 @@ reddit_client_secret = config_data['reddit']['client_secret']
 
 storage_chat = config_data['config'].get('storage_chat')
 storage_msg_id = config_data['config'].get('storage_message_id')
-send_to_chats = config_data['config']['send_to_chats']
-subreddits = config_data['config']['subreddits']
+_bkup_subreddits = config_data['config'].get('subreddits')
+_send_to_chats = config_data['config']['send_to_chats']
+send_to_chats = dict()
+for chat in _send_to_chats:
+    subs = None
+    if isinstance(chat, dict):
+        subs = tuple(chat.values())[0]
+        chat = tuple(chat.keys())[0]
+    send_to_chats[chat] = subs
 bot_admins = config_data['config']['bot_admins']
 cron_duration = config_data['config']['cron_duration']
 
@@ -49,9 +56,12 @@ async def main():
             await (await client.get_messages(storage_chat, ids=storage_msg_id)).download_media('redditbot.json')
         with open('redditbot.json') as file:
             seen_posts = json.load(file)
+        if isinstance(seen_posts, list):
+            seen_posts = {'version': 0, 'chats': {'global': seen_posts}}
     except Exception:
-        traceback.print_exc()
-        seen_posts = []
+        logging.exception('Loading JSON')
+        seen_posts = {'version': 0, 'chats': {'global': []}}
+        # chat dict: {chatid: [array of submission ids]}
 
     async def write_seen_posts():
         with open('redditbot.json', 'w') as file:
@@ -61,47 +71,53 @@ async def main():
 
     @aiocron.crontab(cron_duration)
     async def start_post():
-        while True:
-            random.shuffle(subreddits)
-            to_break = False
-            for subreddit_name in subreddits:
-                subreddit = reddit.subreddit(subreddit_name)
-                while True:
-                    random_post = subreddit.random()
-                    cpid = None
-                    if random_post is None:
-                        for submission in subreddit.hot():
-                            cpid = getattr(submission, 'crosspost_parent', None)
-                            if cpid:
-                                cpid = cpid[3:]
-                            if submission.id in seen_posts or cpid in seen_posts:
-                                continue
-                            random_post = submission
-                            break
-                    cpid = getattr(random_post, 'crosspost_parent', None)
-                    if cpid:
-                        cpid = cpid[3:]
-                    if random_post.id in seen_posts or cpid in seen_posts:
-                        continue
-                    seen_posts.append(cpid or random_post.id)
-                    print(random_post.id, random_post.shortlink)
-                    to_break = True
-                    break
-                if to_break:
-                    break
-            to_break = False
-            for _ in range(5):
+        global_sp = seen_posts['chats']['global']
+        for chat in send_to_chats:
+            uses_bkupsub = False
+            subreddits = send_to_chats[chat]
+            if not subreddits:
+                subreddits = _bkup_subreddits
+                uses_bkupsub = True
+            chat = await client.get_peer_id(chat)
+            if str(chat) not in seen_posts['chats']:
+                seen_posts['chats'][str(chat)] = []
+            chat_sp = global_sp if uses_bkupsub else seen_posts['chats'][str(chat)]
+            while True:
+                random.shuffle(subreddits)
+                to_break = False
+                for subreddit_name in subreddits:
+                    subreddit = reddit.subreddit(subreddit_name)
+                    while True:
+                        random_post = subreddit.random()
+                        cpid = None
+                        if random_post is None:
+                            for submission in subreddit.hot():
+                                cpid = getattr(submission, 'crosspost_parent', None)
+                                if cpid:
+                                    cpid = cpid[3:]
+                                if submission.id in chat_sp + global_sp or cpid in chat_sp + global_sp:
+                                    continue
+                                random_post = submission
+                                break
+                        cpid = getattr(random_post, 'crosspost_parent', None)
+                        if cpid:
+                            cpid = cpid[3:]
+                        if random_post.id in chat_sp + global_sp or cpid in chat_sp + global_sp:
+                            continue
+                        chat_sp.append(cpid or random_post.id)
+                        print(random_post.id, random_post.shortlink)
+                        to_break = True
+                        break
+                    if to_break:
+                        break
                 try:
-                    await _actual_start_post(random_post, send_to_chats)
+                    await _actual_start_post(random_post, [chat])
                 except Exception:
-                    traceback.print_exc()
+                    logging.exception(random_post.id)
                     for i in bot_admins:
                         await client.send_message(i, f'{random_post.id}\n{traceback.format_exc()}', parse_mode=None)
                 else:
-                    to_break = True
                     break
-            if to_break:
-                break
         await write_seen_posts()
 
     async def _start_broadcast(text, file, chats):