Initial commit

main.py

@@ -0,0 +1,124 @@
+import argparse
+import asyncio
+import datetime
+import logging
+
+import discord
+from rich.console import Console
+from rich.logging import RichHandler
+
+from bot.backlog import scan_channel_backlog
+from bot.channels import get_scannable_channels
+from bot.config import load_config
+from bot.discovery import Target, discover_targets
+from bot.display import StatsDisplay
+from bot.reactor import react_to_message
+from bot.storage import Storage
+from bot.undo import undo_all_reactions
+
+log = logging.getLogger("dc_reactor")
+
+
+class ReactorClient(discord.Client):
+    def __init__(self, cfg, storage: Storage, display: StatsDisplay, undo: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        self.cfg = cfg
+        self.storage = storage
+        self.display = display
+        self.undo = undo
+        self.known_channel_targets: dict[int, Target] = {}
+        self._started = False
+
+    async def on_ready(self):
+        log.info("Logged in as %s (%s)", self.user, self.user.id)
+        if self._started:
+            return
+        self._started = True
+        if self.undo:
+            self.loop.create_task(self._run_undo())
+        else:
+            self.loop.create_task(self._run_backlog_then_watch())
+
+    async def _run_undo(self):
+        await undo_all_reactions(self, self.storage, display=self.display)
+        await self.close()
+
+    async def _run_backlog_then_watch(self):
+        targets = await discover_targets(self, self.cfg.friend_id, self.cfg.manual_guild_ids)
+        log.info("Discovered %d mutual target(s): %s", len(targets), [t.name for t in targets])
+
+        for t in targets:
+            await asyncio.to_thread(self.storage.upsert_target, t.target_id, t.target_type, t.name)
+        await self.display.refresh()
+
+        cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=self.cfg.days_back)
+
+        async def react_fn(message: discord.Message, target: Target) -> bool:
+            return await react_to_message(
+                self.storage, message, target, self.cfg.friend_id, self.cfg.emoji,
+                source="backlog", display=self.display,
+            )
+
+        for target in targets:
+            if target.target_type == "group_dm":
+                channel = self.get_channel(target.target_id)
+                if channel is None:
+                    log.warning("Group DM %s not found in cache, skipping", target.target_id)
+                    continue
+                self.known_channel_targets[channel.id] = target
+                await scan_channel_backlog(self.storage, channel, target, self.cfg.friend_id, cutoff, react_fn)
+            else:
+                for channel in get_scannable_channels(self, target.target_id):
+                    self.known_channel_targets[channel.id] = target
+                    await scan_channel_backlog(self.storage, channel, target, self.cfg.friend_id, cutoff, react_fn)
+
+        log.info("Backlog scan complete. Now watching live.")
+
+    async def on_message(self, message: discord.Message):
+        if self.undo:
+            return
+        target = self.known_channel_targets.get(message.channel.id)
+        if target is None:
+            return
+        if message.author.id != self.cfg.friend_id:
+            return
+        await react_to_message(
+            self.storage, message, target, self.cfg.friend_id, self.cfg.emoji,
+            source="live", display=self.display,
+        )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="React to a friend's messages across mutual servers/groups.")
+    parser.add_argument(
+        "--undo",
+        action="store_true",
+        help="Remove every reaction previously added by this script (per the local database log) and exit.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    cfg = load_config()
+
+    console = Console()
+    logging.basicConfig(
+        level=cfg.log_level,
+        format="%(message)s",
+        handlers=[RichHandler(console=console, show_path=False, rich_tracebacks=True)],
+    )
+
+    storage = Storage(cfg.db_path)
+    display = StatsDisplay(console, storage)
+
+    with display:
+        client = ReactorClient(cfg, storage, display, undo=args.undo)
+        try:
+            client.run(cfg.token, log_handler=None)
+        finally:
+            storage.close()
+
+
+if __name__ == "__main__":
+    main()