315 lines
12 KiB
Python
315 lines
12 KiB
Python
# 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:
|
|
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
|