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, resolve_friend 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.known_guild_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): self.display.set_status(f"Resolving friend id {self.cfg.friend_id}...") try: friend = await resolve_friend(self, self.cfg.friend_id) except discord.NotFound: log.error( "friend_id %s does not resolve to any Discord user. Check config.json and fix it.", self.cfg.friend_id, ) self.display.set_status("Failed: friend_id does not resolve to a Discord user.") return log.info("Targeting friend: %s (%s)", friend, friend.id) self.display.set_friend(str(friend)) self.display.set_status(f"Discovering mutual servers/groups with {friend}...") await self.display.refresh() targets = await discover_targets(self, friend, 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) pruned = await asyncio.to_thread(self.storage.prune_targets, [t.target_id for t in targets]) if pruned: log.info("Pruned %d stale target(s) no longer mutual/discovered", pruned) self.display.set_status(f"Starting backlog scan across {len(targets)} target(s)...") 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, count_found=False, ) for target in targets: self.display.set_active_target(target.target_id) if target.target_type in ("group_dm", "dm"): channel = self.get_channel(target.target_id) if channel is None: log.warning("Channel %s (%s) not found in cache, skipping", target.target_id, target.name) continue self.known_channel_targets[channel.id] = target await scan_channel_backlog( self.storage, channel, target, self.cfg.friend_id, self.cfg.emoji, cutoff, react_fn, display=self.display, ) else: self.known_guild_targets[target.target_id] = target 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, self.cfg.emoji, cutoff, react_fn, display=self.display, ) self.display.set_active_target(None) log.info("Backlog scan complete. Now watching live.") self.display.set_status("Backlog scan complete — watching live for new messages...") 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: # Channel didn't exist (or wasn't enumerated, e.g. a thread) at # backlog-scan time. If it belongs to an already-known mutual # guild, start watching it now instead of requiring a restart — # only its live messages going forward are covered, no backfill. guild = getattr(message, "guild", None) if guild is not None: target = self.known_guild_targets.get(guild.id) if target is not None: self.known_channel_targets[message.channel.id] = target log.info( "New channel #%s in %s wasn't seen during backlog scan; watching it live from now on", getattr(message.channel, "name", message.channel.id), target.name, ) if target is None: return if message.author.id != self.cfg.friend_id: return chan_label = getattr(message.channel, "name", None) or target.name self.display.set_active_target(target.target_id) self.display.set_status(f"New message from friend in {chan_label}...") await react_to_message( self.storage, message, target, self.cfg.friend_id, self.cfg.emoji, source="live", display=self.display, ) self.display.set_active_target(None) 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) display.set_status("Connecting to Discord...") 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()