# 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 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. """ 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 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!" 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"SUCCESS: {STATE.friendly[friendly_id].name} successfully hospitalized {enemy.name}") # Remove from active targets del self.active_targets[key] continue # Send reminder 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 await discord_user.send(f"**Reminder:** 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: 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