Basic Discord functionality with assignments
This commit is contained in:
246
services/bot_assignment.py
Normal file
246
services/bot_assignment.py
Normal file
@@ -0,0 +1,246 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user