# services/bot_assignment.py import asyncio import json from pathlib import Path from typing import Dict, Optional from datetime import datetime from services.server_state import STATE from config import ASSIGNMENT_TIMEOUT, ASSIGNMENT_REMINDER, ALLOWED_CHANNEL_ID class BotAssignmentManager: def __init__(self, bot): self.bot = bot self.discord_mapping: Dict[int, int] = {} # torn_id -> discord_id self.active_targets: Dict[str, Dict] = {} # key: "group_id:enemy_id" -> assignment data self.running = False self.task = None # Load Discord mapping self.load_discord_mapping() def load_discord_mapping(self): """Load Torn ID to Discord ID mapping from JSON file""" path = Path("data/discord_mapping.json") if not path.exists(): print("No discord_mapping.json found") return try: with open(path, "r", encoding="utf-8") as f: data = json.load(f) # Convert string keys to int self.discord_mapping = {int(k): int(v) for k, v in data.get("mappings", {}).items()} print(f"Loaded {len(self.discord_mapping)} Discord mappings") except Exception as e: print(f"Error loading discord mapping: {e}") def get_discord_id(self, torn_id: int) -> Optional[int]: """Get Discord user ID for a Torn player ID""" return self.discord_mapping.get(torn_id) async def start(self): """Start the bot assignment loop""" if self.running: print("WARNING: Bot assignment already running") return self.running = True self.task = asyncio.create_task(self.assignment_loop()) print("Bot assignment loop started") print(f"Loaded {len(self.discord_mapping)} Discord ID mappings") async def stop(self): """Stop the bot assignment loop""" if not self.running: return self.running = False if self.task: self.task.cancel() try: await self.task except asyncio.CancelledError: pass print("Bot assignment stopped") def friendly_has_active_target(self, friendly_id: int) -> bool: """Check if a friendly player already has an active target assigned""" for target_data in self.active_targets.values(): if target_data["friendly_id"] == friendly_id: return True return False def get_next_friendly_in_group(self, group_id: str, friendly_ids: list) -> Optional[int]: """ Get the next friendly in the group who should receive a target. Prioritizes members with fewer hits. Only returns friendlies who DON'T already have an active assignment. """ if not friendly_ids: return None # Get hit counts for all friendlies in this group who DON'T have active targets friendly_hits = [] for fid in friendly_ids: if fid in STATE.friendly: # Skip if this friendly already has an active target if self.friendly_has_active_target(fid): continue hits = STATE.friendly[fid].hits friendly_hits.append((fid, hits)) if not friendly_hits: return None # Sort by hit count (ascending) - members with fewer hits first friendly_hits.sort(key=lambda x: x[1]) # Return the friendly with the fewest hits who doesn't have an active target return friendly_hits[0][0] def get_next_enemy_in_group(self, group_id: str, enemy_ids: list) -> Optional[int]: """ Get the next enemy in the group who needs to be assigned. Returns None if all enemies are already assigned or not attackable. """ for eid in enemy_ids: key = f"{group_id}:{eid}" # If enemy is already assigned, skip them if key in self.active_targets: continue # Check if enemy is attackable (status must be "Okay") enemy = STATE.enemy.get(eid) if not enemy or enemy.status.lower() != "okay": continue # This enemy is available for assignment return eid return None async def assignment_loop(self): """Main loop that assigns targets and monitors status""" await self.bot.wait_until_ready() print("Bot is ready, assignment loop running") first_run = True while self.running: try: # Check if bot is enabled via STATE if not STATE.bot_running: if first_run: print("Bot paused - waiting for Start Bot button to be clicked") first_run = False await asyncio.sleep(5) continue if first_run: print("Bot activated - processing assignments") first_run = False # Process each group has_assignments = False async with STATE.lock: for group_id, assignments in STATE.groups.items(): friendly_ids = assignments.get("friendly", []) enemy_ids = assignments.get("enemy", []) if friendly_ids and enemy_ids: has_assignments = True if not friendly_ids or not enemy_ids: continue # Try to assign any unassigned enemies enemy_id = self.get_next_enemy_in_group(group_id, enemy_ids) if enemy_id: friendly_id = self.get_next_friendly_in_group(group_id, friendly_ids) if friendly_id: await self.assign_target(group_id, friendly_id, enemy_id) if not has_assignments and STATE.bot_running: print("No group assignments found - drag members into groups first!") # Monitor active targets for status changes or timeouts await self.monitor_active_targets() # Sleep before next iteration await asyncio.sleep(5) except Exception as e: print(f"Error in assignment loop: {e}") import traceback traceback.print_exc() await asyncio.sleep(5) async def assign_target(self, group_id: str, friendly_id: int, enemy_id: int): """Assign an enemy target to a friendly player""" # Get member data friendly = STATE.friendly.get(friendly_id) enemy = STATE.enemy.get(enemy_id) if not friendly or not enemy: print(f"Cannot assign: friendly {friendly_id} or enemy {enemy_id} not found") return # Only assign if enemy status is "Okay" (not traveling, hospitalized, etc.) if enemy.status.lower() != "okay": print(f"Skipping assignment: {enemy.name} status is '{enemy.status}' (must be 'Okay')") return # Get Discord user discord_id = self.get_discord_id(friendly_id) if not discord_id: print(f"No Discord mapping for Torn ID {friendly_id} - skipping assignment") # Record assignment to prevent infinite retries key = f"{group_id}:{enemy_id}" self.active_targets[key] = { "group_id": group_id, "friendly_id": friendly_id, "enemy_id": enemy_id, "discord_id": None, "assigned_at": datetime.now(), "reminded": False, "failed": True } return # Record assignment BEFORE attempting to send message # This prevents infinite retries if message sending fails key = f"{group_id}:{enemy_id}" self.active_targets[key] = { "group_id": group_id, "friendly_id": friendly_id, "enemy_id": enemy_id, "discord_id": discord_id, "assigned_at": datetime.now(), "reminded": False, "failed": False } # Fetch Discord user and send message try: discord_user = await self.bot.fetch_user(discord_id) if not discord_user: print(f"Discord user {discord_id} not found") self.active_targets[key]["failed"] = True return # Send Discord message to channel attack_link = f"https://www.torn.com/loader.php?sid=attack&user2ID={enemy_id}" message = f"**New target for {discord_user.mention}!**\n\n[**{enemy.name}** (Level {enemy.level})]({attack_link})\n\nYou have {ASSIGNMENT_TIMEOUT} seconds!" channel = self.bot.get_channel(ALLOWED_CHANNEL_ID) if channel: await channel.send(message) print(f"Assigned {enemy.name} to {friendly.name} (Discord: {discord_user.name})") else: print(f"Assignment channel {ALLOWED_CHANNEL_ID} not found") self.active_targets[key]["failed"] = True except Exception as e: print(f"Failed to send Discord message to channel: {e}") self.active_targets[key]["failed"] = True async def monitor_active_targets(self): """Monitor active targets for status changes or timeouts""" now = datetime.now() to_reassign = [] for key, data in list(self.active_targets.items()): elapsed = (now - data["assigned_at"]).total_seconds() # Remove failed assignments after a short delay (don't reassign them) if data.get("failed", False): if elapsed >= 30: # Clean up after 30 seconds print(f"Removing failed assignment: {key}") del self.active_targets[key] continue # Check enemy status enemy_id = data["enemy_id"] enemy = STATE.enemy.get(enemy_id) if not enemy: # Enemy no longer exists, remove assignment del self.active_targets[key] continue # Check if enemy is hospitalized (success!) if "hospital" in enemy.status.lower(): friendly_id = data["friendly_id"] if friendly_id in STATE.friendly: # Increment hit count STATE.friendly[friendly_id].hits += 1 print(f"SUCCESS: {STATE.friendly[friendly_id].name} successfully hospitalized {enemy.name}") # Remove from active targets del self.active_targets[key] continue # Check if enemy is no longer attackable (traveling, etc.) if enemy.status.lower() != "okay": print(f"Target {enemy.name} is now '{enemy.status}' - removing assignment") del self.active_targets[key] continue # Send reminder (only for successful assignments) if elapsed >= ASSIGNMENT_REMINDER and not data["reminded"]: discord_id = data["discord_id"] try: discord_user = await self.bot.fetch_user(discord_id) remaining = ASSIGNMENT_TIMEOUT - ASSIGNMENT_REMINDER channel = self.bot.get_channel(ALLOWED_CHANNEL_ID) if channel: await channel.send(f"**Reminder:** {discord_user.mention} - Target {enemy.name} - {remaining} seconds left!") data["reminded"] = True except: pass # Reassign after timeout if elapsed >= ASSIGNMENT_TIMEOUT: to_reassign.append((data["group_id"], enemy_id)) del self.active_targets[key] # Reassign targets that timed out for group_id, enemy_id in to_reassign: # Verify enemy is still in this group before reassigning enemy_ids = STATE.groups[group_id].get("enemy", []) if enemy_id not in enemy_ids: print(f"Enemy {enemy_id} no longer in group {group_id} - not reassigning") continue friendly_ids = STATE.groups[group_id].get("friendly", []) friendly_id = self.get_next_friendly_in_group(group_id, friendly_ids) if friendly_id: print(f"Reassigning enemy {enemy_id} (timeout)") await self.assign_target(group_id, friendly_id, enemy_id) # Global instance (will be initialized with bot in main.py) assignment_manager: Optional[BotAssignmentManager] = None