# 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 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("⚠ 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 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. """ if not friendly_ids: return None # Get hit counts for all friendlies in this group friendly_hits = [] for fid in friendly_ids: if fid in STATE.friendly: 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 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. """ for eid in enemy_ids: key = f"{group_id}:{eid}" # If enemy is not currently assigned, return it if key not in self.active_targets: 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 # 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}") return discord_user = await self.bot.fetch_user(discord_id) if not discord_user: print(f"Discord user {discord_id} not found") return # Record assignment 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 } # Send Discord message enemy_link = f"https://www.torn.com/profiles.php?XID={enemy_id}" message = f"🎯 **New target for {discord_user.mention}!**\n\nAttack **{enemy.name}** (Level {enemy.level})\n{enemy_link}\n\n⏰ You have 30 seconds!" try: await discord_user.send(message) print(f"Assigned {enemy.name} to {friendly.name} (Discord: {discord_user.name})") except Exception as e: print(f"Failed to send Discord message to {discord_user.name}: {e}") 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() # Check enemy status enemy_id = data["enemy_id"] enemy = STATE.enemy.get(enemy_id) if enemy and "hospital" in enemy.status.lower(): # Enemy is hospitalized - success! friendly_id = data["friendly_id"] if friendly_id in STATE.friendly: # Increment hit count STATE.friendly[friendly_id].hits += 1 print(f"✓ {STATE.friendly[friendly_id].name} successfully hospitalized {enemy.name}") # Remove from active targets del self.active_targets[key] continue # Send reminder at 15 seconds if elapsed >= 15 and not data["reminded"]: discord_id = data["discord_id"] try: discord_user = await self.bot.fetch_user(discord_id) await discord_user.send(f"⏰ **Reminder:** Target {enemy.name} - 15 seconds left!") data["reminded"] = True except: pass # Reassign after 30 seconds if elapsed >= 30: 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: 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