259 lines
9.8 KiB
Python
259 lines
9.8 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
|
|
|
|
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 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
|
|
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 60 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 60 seconds
|
|
if elapsed >= 60:
|
|
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
|