faction_war_dispatch_bot/services/bot_assignment.py

# services/bot_assignment.py
import asyncio
import json
import aiohttp
from pathlib import Path
from typing import Dict, Optional
from datetime import datetime
from services.server_state import STATE
from services.activity_log import activity_logger
import config

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

        # Chain monitoring state
        self.chain_timeout = 0  # Seconds remaining on chain
        self.chain_active = False  # Whether we're in assignment mode
        self.assigned_friendlies = set()  # Track who has been assigned in this chain session
        self.current_group_index = 0  # For round-robin group assignment
        self.chain_warning_sent = False  # Track if 30-second warning was sent
        self.last_chain_check = datetime.now()

        # 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 fetch_chain_timer(self) -> Optional[int]:
        #Fetch chain timeout from Torn API. Returns seconds remaining, or None if no chain or error.
        if not STATE.friendly_faction_id:
            return None

        try:
            url = f"https://api.torn.com/v2/faction/{STATE.friendly_faction_id}/chain?key={config.TORN_API_KEY}"
            async with aiohttp.ClientSession() as session:
                async with session.get(url) as resp:
                    if resp.status != 200:
                        return None
                    data = await resp.json()

            # Get timeout value in seconds
            chain_data = data.get("chain", {})
            timeout = chain_data.get("timeout", 0)
            return timeout
        except Exception as e:
            print(f"Error fetching chain timer: {e}")
            return None

    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 check_chain_timer(self):
        #Check chain timer and update chain state
        timeout = await self.fetch_chain_timer()

        if timeout is None:
            # No chain data or error
            if self.chain_active:
                print("Chain ended or API error - resetting chain state")
                self.chain_active = False
                self.assigned_friendlies.clear()
                self.current_group_index = 0
                self.chain_warning_sent = False
            return

        self.chain_timeout = timeout
        threshold_seconds = config.CHAIN_TIMER_THRESHOLD * 60

        # Check if we should enter chain mode
        if timeout > 0 and timeout <= threshold_seconds and not self.chain_active:
            print(f"Chain timer at {timeout}s (threshold: {threshold_seconds}s) - entering chain mode")
            self.chain_active = True
            self.assigned_friendlies.clear()
            self.current_group_index = 0
            self.chain_warning_sent = False
            # Log to activity
            await activity_logger.log_action("System", "Chain Mode Activated", f"Timer: {timeout}s, Threshold: {threshold_seconds}s")

        # Check if chain expired
        elif timeout > threshold_seconds and self.chain_active:
            print(f"Chain timer above threshold ({timeout}s > {threshold_seconds}s) - exiting chain mode")
            self.chain_active = False
            self.assigned_friendlies.clear()
            self.current_group_index = 0
            self.chain_warning_sent = False
            # Log to activity
            await activity_logger.log_action("System", "Chain Mode Deactivated", f"Timer: {timeout}s exceeded threshold")

        # Check if chain expired (timeout = 0)
        elif timeout == 0 and self.chain_active:
            print("Chain expired - resetting chain state")
            self.chain_active = False
            self.assigned_friendlies.clear()
            self.current_group_index = 0
            self.chain_warning_sent = False
            # Log to activity
            await activity_logger.log_action("System", "Chain Expired", "Chain timer reached 0")

        # Send 30-second warning
        if self.chain_active and timeout <= 30 and not self.chain_warning_sent:
            await self.send_chain_expiration_warning()
            self.chain_warning_sent = True
            # Log to activity
            await activity_logger.log_action("System", "Chain Expiration Warning", "30 seconds remaining")

    async def send_chain_expiration_warning(self):
        #Send @here alert that chain is about to expire
        try:
            channel = self.bot.get_channel(config.ALLOWED_CHANNEL_ID)
            if channel and STATE.friendly_faction_id:
                faction_link = f"https://www.torn.com/factions.php?step=your#/tab=chain"
                message = f"@here **CHAIN EXPIRING IN 30 SECONDS!** Attack a faction member to keep it alive!\n{faction_link}"
                await channel.send(message)
                print("Sent chain expiration warning")
        except Exception as e:
            print(f"Error sending chain warning: {e}")

    async def assign_next_chain_hit(self):
        #Assign next hit in round-robin fashion through groups
        # Only assign if there's no active assignment waiting
        if self.active_targets:
            return  # Wait for current assignment to complete

        # Get list of group IDs (sorted for consistency)
        group_ids = sorted(STATE.groups.keys())
        if not group_ids:
            return

        # Try to find a group with available friendly and enemy
        attempts = 0
        while attempts < len(group_ids):
            # Get current group
            group_id = group_ids[self.current_group_index % len(group_ids)]

            async with STATE.lock:
                friendly_ids = STATE.groups[group_id].get("friendly", [])
                enemy_ids = STATE.groups[group_id].get("enemy", [])

                # Find a friendly who hasn't been assigned yet
                available_friendly = None
                for fid in friendly_ids:
                    if fid not in self.assigned_friendlies:
                        available_friendly = fid
                        break

                # Find an attackable enemy
                enemy_id = self.get_next_enemy_in_group(group_id, enemy_ids)

                if available_friendly and enemy_id:
                    # Make assignment
                    self.assigned_friendlies.add(available_friendly)
                    await self.assign_target(group_id, available_friendly, enemy_id)
                    print(f"Chain hit assigned: Group {group_id}, Friendly {available_friendly} -> Enemy {enemy_id}")

                    # Move to next group for next assignment
                    self.current_group_index += 1
                    return

            # Move to next group and try again
            self.current_group_index += 1
            attempts += 1

        # If we've tried all groups and no assignments possible
        # Check if we need to reset assigned_friendlies (everyone has been assigned)
        async with STATE.lock:
            all_friendlies = set()
            for assignments in STATE.groups.values():
                all_friendlies.update(assignments.get("friendly", []))

            if self.assigned_friendlies and self.assigned_friendlies >= all_friendlies:
                print("All friendlies have been assigned - resetting for round-robin")
                self.assigned_friendlies.clear()
                self.current_group_index = 0

    async def assignment_loop(self):
        #Main loop that monitors chain timer and assigns targets
        await self.bot.wait_until_ready()
        print("Bot is ready, assignment loop running with chain timer monitoring")

        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 - chain timer monitoring enabled")
                    first_run = False

                # Check chain timer every 30 seconds
                now = datetime.now()
                if (now - self.last_chain_check).total_seconds() >= 30:
                    await self.check_chain_timer()
                    self.last_chain_check = now

                # If chain is active, try to assign next hit
                if self.chain_active:
                    await self.assign_next_chain_hit()

                # 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 {config.ASSIGNMENT_TIMEOUT} seconds!"

            channel = self.bot.get_channel(config.ALLOWED_CHANNEL_ID)
            if channel:
                await channel.send(message)
                print(f"Assigned {enemy.name} to {friendly.name} (Discord: {discord_user.name})")
                # Log to activity
                await activity_logger.log_action("System", "Hit Assigned", f"{friendly.name} -> {enemy.name} (Level {enemy.level})")
            else:
                print(f"Assignment channel {config.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
                    friendly_name = STATE.friendly[friendly_id].name
                    enemy_name = enemy.name
                    print(f"SUCCESS: {friendly_name} successfully hospitalized {enemy_name}")
                    # Log to activity
                    await activity_logger.log_action("System", "Hit Completed", f"{friendly_name} 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":
                friendly_id = data["friendly_id"]
                friendly_name = STATE.friendly.get(friendly_id).name if friendly_id in STATE.friendly else "Unknown"
                enemy_name = enemy.name
                enemy_status = enemy.status
                print(f"Target {enemy_name} is now '{enemy_status}' - removing assignment")
                # Log to activity
                await activity_logger.log_action("System", "Hit Dropped", f"{friendly_name}'s target {enemy_name} became {enemy_status}")
                del self.active_targets[key]
                continue

            # Send reminder (only for successful assignments)
            if elapsed >= config.ASSIGNMENT_REMINDER and not data["reminded"]:
                discord_id = data["discord_id"]
                try:
                    discord_user = await self.bot.fetch_user(discord_id)
                    remaining = config.ASSIGNMENT_TIMEOUT - config.ASSIGNMENT_REMINDER
                    channel = self.bot.get_channel(config.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 >= config.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:
            # Verify enemy is still in this group before reassigning
            enemy_ids = STATE.groups[group_id].get("enemy", [])
            if enemy_id not in enemy_ids:
                print(f"Enemy {enemy_id} no longer in group {group_id} - not reassigning")
                continue

            friendly_ids = STATE.groups[group_id].get("friendly", [])
            friendly_id = self.get_next_friendly_in_group(group_id, friendly_ids)
            if friendly_id:
                enemy_name = STATE.enemy.get(enemy_id).name if enemy_id in STATE.enemy else f"Enemy {enemy_id}"
                friendly_name = STATE.friendly.get(friendly_id).name if friendly_id in STATE.friendly else f"Friendly {friendly_id}"
                print(f"Reassigning enemy {enemy_id} (timeout)")
                # Log to activity
                await activity_logger.log_action("System", "Hit Timed Out", f"Reassigning {enemy_name} to {friendly_name} (previous assignee timed out)")
                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