#!/usr/bin/env python3
"""
CLI tool to enforce policies on profiles.
"""
import argparse
import json
import logging
import os
import signal
import sys
import time

try:
    import yaml
except ImportError:
    print("PyYAML is not installed. Please install it with: pip install PyYAML", file=sys.stderr)
    yaml = None

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = None

from .profile_manager_tool import ProfileManager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Graceful shutdown handler
shutdown_event = False
def handle_shutdown(sig, frame):
    global shutdown_event
    logger.info("Shutdown signal received. Stopping policy enforcer...")
    shutdown_event = True

class PolicyEnforcer:
    def __init__(self, manager, dry_run=False):
        self.manager = manager
        self.dry_run = dry_run
        self.actions_taken_this_cycle = 0

    PROXY_REST_REASON = "Proxy resting"

    def apply_policies(self, args):
        self.actions_taken_this_cycle = 0
        logger.debug(f"Applying policies... (Dry run: {self.dry_run})")
        
        # --- Phase 1: Policies that don't depend on a consistent profile state snapshot ---
        # Manage proxy states and clean up stale locks before we fetch profile states.
        self.enforce_proxy_group_rotation(getattr(args, 'proxy_groups', []))
        self.enforce_proxy_work_rest_cycle(args)
        self.enforce_max_proxy_active_time(args)
        self.enforce_proxy_policies(args)
        if args.unlock_stale_locks_after_seconds and args.unlock_stale_locks_after_seconds > 0:
            self.enforce_stale_lock_cleanup(args.unlock_stale_locks_after_seconds)

        # --- Phase 2: Policies that require a consistent, shared view of profile states ---
        # Fetch all profile states ONCE to create a consistent snapshot for this cycle.
        all_profiles_list = self.manager.list_profiles()
        all_profiles_map = {p['name']: p for p in all_profiles_list}

        # Apply profile group policies (rotation, max_active). This will modify the local `all_profiles_map`.
        self.enforce_profile_group_policies(getattr(args, 'profile_groups', []), all_profiles_map)

        # Un-rest profiles. This also reads from and modifies the local `all_profiles_map`.
        self.enforce_unrest_policy(getattr(args, 'profile_groups', []), all_profiles_map)
        
        # --- Phase 3: Apply policies to individual active profiles ---
        # Use the now-updated snapshot to determine which profiles are active.
        active_profiles = [p for p in all_profiles_map.values() if p['state'] == self.manager.STATE_ACTIVE]

        # Filter out profiles that are managed by a profile group, as their state is handled separately.
        profile_groups = getattr(args, 'profile_groups', [])
        if profile_groups:
            grouped_profiles = set()
            for group in profile_groups:
                if 'profiles' in group:
                    for p_name in group['profiles']:
                        grouped_profiles.add(p_name)
                elif 'prefix' in group:
                    prefix = group['prefix']
                    for p in all_profiles_list:
                        if p['name'].startswith(prefix):
                            grouped_profiles.add(p['name'])
            
            original_count = len(active_profiles)
            active_profiles = [p for p in active_profiles if p['name'] not in grouped_profiles]
            if len(active_profiles) != original_count:
                logger.debug(f"Filtered out {original_count - len(active_profiles)} profile(s) managed by profile groups.")

        for profile in active_profiles:
            # Check for failure burst first, as it's more severe.
            # If it's banned, no need to check other rules for it.
            if self.enforce_failure_burst_policy(profile, args.ban_on_failures, args.ban_on_failures_window_minutes):
                continue
            if self.enforce_rate_limit_policy(profile, getattr(args, 'rate_limit_requests', 0), getattr(args, 'rate_limit_window_minutes', 0), getattr(args, 'rate_limit_rest_duration_minutes', 0)):
                continue
            self.enforce_failure_rate_policy(profile, args.max_failure_rate, args.min_requests_for_rate)
            self.enforce_rest_policy(profile, args.rest_after_requests, args.rest_duration_minutes)

        return self.actions_taken_this_cycle > 0

    def enforce_failure_burst_policy(self, profile, max_failures, window_minutes):
        if not max_failures or not window_minutes or max_failures <= 0 or window_minutes <= 0:
            return False
        
        window_seconds = window_minutes * 60
        # Count only fatal error types (auth, download) for the ban policy.
        # Tolerated errors are excluded from this check.
        error_count = (
            self.manager.get_activity_rate(profile['name'], 'failure', window_seconds) +
            self.manager.get_activity_rate(profile['name'], 'download_error', window_seconds)
        )
        
        if error_count >= max_failures:
            reason = f"Error burst detected: {error_count} errors in the last {window_minutes} minute(s) (threshold: {max_failures})"
            logger.warning(f"Banning profile '{profile['name']}' due to error burst: {reason}")
            if not self.dry_run:
                self.manager.update_profile_state(profile['name'], self.manager.STATE_BANNED, reason)
            self.actions_taken_this_cycle += 1
            return True # Indicates profile was banned
        return False

    def enforce_rate_limit_policy(self, profile, max_requests, window_minutes, rest_duration_minutes):
        if not max_requests or not window_minutes or max_requests <= 0 or window_minutes <= 0:
            return False

        window_seconds = window_minutes * 60
        # Count all successful activities (auth, download) for rate limiting.
        # We don't count failures, as they often don't hit the target server in the same way.
        activity_count = (
            self.manager.get_activity_rate(profile['name'], 'success', window_seconds) +
            self.manager.get_activity_rate(profile['name'], 'download', window_seconds)
        )

        if activity_count >= max_requests:
            reason = f"Rate limit hit: {activity_count} requests in last {window_minutes} minute(s) (limit: {max_requests})"
            logger.info(f"Resting profile '{profile['name']}' for {rest_duration_minutes}m: {reason}")
            if not self.dry_run:
                self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, reason)
                rest_until = time.time() + rest_duration_minutes * 60
                self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until))
            self.actions_taken_this_cycle += 1
            return True # Indicates profile was rested
        return False

    def enforce_unrest_policy(self, profile_groups, all_profiles_map):
        all_profiles_list = list(all_profiles_map.values())
        resting_profiles = [p for p in all_profiles_list if p['state'] == self.manager.STATE_RESTING]
        cooldown_profiles = [p for p in all_profiles_list if p['state'] == self.manager.STATE_COOLDOWN]
        profiles_to_check = resting_profiles + cooldown_profiles
        now = time.time()

        if not profiles_to_check:
            return

        # Sort profiles to check by their rest_until timestamp, then by name.
        # This creates a deterministic FIFO queue for activation.
        profiles_to_check.sort(key=lambda p: (p.get('rest_until', 0), p.get('name', '')))

        # --- Group-aware unrest logic ---
        profile_to_group_map = {}
        group_to_profiles_map = {}
        if profile_groups:
            for group in profile_groups:
                group_name = group.get('name')
                if not group_name: continue
                
                profiles_in_group = []
                if 'profiles' in group:
                    profiles_in_group = sorted(group['profiles'])
                elif 'prefix' in group:
                    prefix = group['prefix']
                    profiles_in_group = sorted([p['name'] for p in all_profiles_list if p['name'].startswith(prefix)])
                
                group_to_profiles_map[group_name] = profiles_in_group
                for p_name in profiles_in_group:
                    profile_to_group_map[p_name] = group_name
        
        # This will store the live count of active profiles for each group,
        # preventing race conditions within a single enforcer run.
        live_active_counts = {}
        if profile_groups:
            for group_name, profiles_in_group in group_to_profiles_map.items():
                count = 0
                for p_name in profiles_in_group:
                    profile_state = all_profiles_map.get(p_name, {}).get('state')
                    if profile_state in [self.manager.STATE_ACTIVE, self.manager.STATE_LOCKED, self.manager.STATE_COOLDOWN]:
                        count += 1
                live_active_counts[group_name] = count
        # --- End group logic setup ---

        # --- New logic: Identify groups with waiting profiles ---
        groups_with_waiting_profiles = {}
        if profile_groups:
            for group in profile_groups:
                group_name = group.get('name')
                if not group_name: continue
                
                defer_activation = group.get('defer_activation_if_any_waiting', False)
                if not defer_activation: continue

                profiles_in_group = group_to_profiles_map.get(group_name, [])
                waiting_profile = next(
                    (p for p_name, p in all_profiles_map.items()
                     if p_name in profiles_in_group and p.get('rest_reason') == 'waiting_downloads'),
                    None
                )
                if waiting_profile:
                    groups_with_waiting_profiles[group_name] = waiting_profile['name']
        # --- End new logic ---

        unique_proxies = sorted(list(set(p['proxy'] for p in profiles_to_check if p.get('proxy'))))
        proxy_states = self.manager.get_proxy_states(unique_proxies)
        
        for profile in profiles_to_check:
            profile_name = profile['name']
            group_name = profile_to_group_map.get(profile_name)

            # --- New logic: Defer activation if group has a waiting profile ---
            if group_name in groups_with_waiting_profiles:
                waiting_profile_name = groups_with_waiting_profiles[group_name]
                if profile_name != waiting_profile_name:
                    logger.debug(f"Profile '{profile_name}' activation deferred because profile '{waiting_profile_name}' in group '{group_name}' is waiting for downloads.")
                    continue
            # --- End new logic ---

            # --- New logic for waiting_downloads ---
            if profile.get('rest_reason') == 'waiting_downloads':
                profile_name = profile['name']
                group_name = profile_to_group_map.get(profile_name)
                group_policy = next((g for g in profile_groups if g.get('name') == group_name), {})
                
                max_wait_minutes = group_policy.get('max_wait_for_downloads_minutes', 240)
                wait_started_at = profile.get('wait_started_at', 0)
                
                downloads_pending = self.manager.get_pending_downloads(profile_name)
                is_timed_out = (time.time() - wait_started_at) > (max_wait_minutes * 60) if wait_started_at > 0 else False

                if downloads_pending <= 0 or is_timed_out:
                    if is_timed_out:
                        logger.warning(f"Profile '{profile_name}' download wait timed out after {max_wait_minutes}m. Forcing rotation.")
                    else:
                        logger.info(f"All pending downloads for profile '{profile_name}' are complete. Proceeding with rotation.")
                    
                    self.actions_taken_this_cycle += 1
                    
                    # Transition to a normal post-rotation rest period.
                    new_reason = "Rotation complete (downloads finished)"
                    rest_duration_minutes = group_policy.get('rest_duration_minutes_on_rotation', 0)
                    rest_until_ts = time.time() + (rest_duration_minutes * 60)

                    if not self.dry_run:
                        self.manager.update_profile_field(profile['name'], 'rest_reason', new_reason)
                        self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
                        self.manager.clear_pending_downloads(profile_name)
                        # Reset counters to ensure a fresh start after the wait period.
                        self.manager.reset_profile_counters(profile_name)
                    
                    # Update local map so it can be activated in the same cycle if rest_duration is 0
                    all_profiles_map[profile_name]['rest_reason'] = new_reason
                    all_profiles_map[profile_name]['rest_until'] = rest_until_ts
                    all_profiles_map[profile_name]['success_count'] = 0
                    all_profiles_map[profile_name]['failure_count'] = 0
                    all_profiles_map[profile_name]['tolerated_error_count'] = 0
                    all_profiles_map[profile_name]['download_count'] = 0
                    all_profiles_map[profile_name]['download_error_count'] = 0
                    
                    # Let the rest of the unrest logic handle the activation now that rest_until is set.
                    profile['rest_until'] = rest_until_ts # Update profile in loop
                else:
                    logger.debug(f"Profile '{profile_name}' is still waiting for {downloads_pending} download(s) to complete.")
                    continue # Skip to next profile, do not attempt to activate.
            # --- End new logic ---

            rest_until = profile.get('rest_until', 0)
            if now >= rest_until:
                profile_name = profile['name']
                group_name = profile_to_group_map.get(profile_name)

                # --- Group-aware unrest check ---
                if group_name:
                    group_policy = next((g for g in profile_groups if g.get('name') == group_name), None)
                    if not group_policy:
                        continue # Should not happen if maps are built correctly

                    max_active = group_policy.get('max_active_profiles', 1)
                    
                    # Check if the group is already at its capacity for active profiles.
                    # We use the live counter which is updated during this enforcer cycle.
                    
                    # Special handling for COOLDOWN profiles: they should be allowed to become ACTIVE
                    # even if the group is at capacity, because they are already counted as "active".
                    # We check if the group would be over capacity *without* this profile.
                    is_cooldown_profile = profile['state'] == self.manager.STATE_COOLDOWN
                    effective_active_count = live_active_counts.get(group_name, 0)
                    
                    # If we are considering a COOLDOWN profile, it's already in the count.
                    # The check should be if activating it would exceed the limit, assuming
                    # no *other* profile is active.
                    capacity_check_count = effective_active_count
                    if is_cooldown_profile:
                        capacity_check_count -= 1

                    if capacity_check_count >= max_active:
                        logger.debug(f"Profile '{profile_name}' rest ended, but group '{group_name}' is at capacity ({effective_active_count}/{max_active}). Deferring activation.")
                        
                        # If a profile's COOLDOWN ends but it can't be activated (because another
                        # profile is active), move it to RESTING so it's clear it's waiting for capacity.
                        if is_cooldown_profile:
                            reason = "Waiting for group capacity"
                            logger.info(f"Profile '{profile_name}' cooldown ended but group is full. Moving to RESTING to wait for a slot.")
                            self.actions_taken_this_cycle += 1
                            if not self.dry_run:
                                # We need to update state but keep rest_until in the past.
                                self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, reason)
                                self.manager.update_profile_field(profile['name'], 'rest_until', '0')
                            
                            # Update local map
                            all_profiles_map[profile_name]['state'] = self.manager.STATE_RESTING
                            all_profiles_map[profile_name]['rest_until'] = 0
                            all_profiles_map[profile_name]['rest_reason'] = reason

                        continue # Do not activate, group is full.
                else:
                    # Defensive check for orphaned profiles that should be in a group.
                    # This can happen if list_profiles() returns an incomplete list for one cycle,
                    # causing the group maps to be incomplete. This check prevents a "stampede"
                    # of activations that would violate group limits.
                    is_orphan = False
                    for group in profile_groups:
                        prefix = group.get('prefix')
                        if prefix and profile_name.startswith(prefix):
                            is_orphan = True
                            logger.warning(
                                f"Profile '{profile_name}' appears to belong to group '{group.get('name')}' "
                                f"but was not found in the initial scan. Deferring activation to prevent violating group limits."
                            )
                            break
                    if is_orphan:
                        continue # Skip activation for this profile
                # --- End group check ---

                # Before activating, ensure the profile's proxy is not resting.
                proxy_url = profile.get('proxy')
                if proxy_url:
                    proxy_state_data = proxy_states.get(proxy_url, {})
                    if proxy_state_data.get('state') == self.manager.STATE_RESTING:
                        logger.debug(f"Profile '{profile['name']}' rest period ended, but its proxy '{proxy_url}' is still resting. Deferring activation.")
                        
                        # Update reason for clarity in the UI when a profile is blocked by its proxy.
                        new_reason = "Waiting for proxy"
                        if profile.get('rest_reason') != new_reason:
                            logger.info(f"Updating profile '{profile['name']}' reason to '{new_reason}'.")
                            self.actions_taken_this_cycle += 1
                            if not self.dry_run:
                                self.manager.update_profile_field(profile['name'], 'rest_reason', new_reason)
                            # Update local map for consistency within this enforcer cycle.
                            all_profiles_map[profile_name]['rest_reason'] = new_reason
                        
                        continue # Do not activate this profile yet.

                # Update group counter BEFORE making any changes, so subsequent checks in this cycle use the updated count
                if group_name and profile['state'] == self.manager.STATE_RESTING:
                    # For RESTING profiles, they're becoming active, so increment the count
                    live_active_counts[group_name] = live_active_counts.get(group_name, 0) + 1
                # COOLDOWN profiles are already counted, no change needed
                
                logger.info(f"Activating profile '{profile['name']}' (rest period completed).")
                self.actions_taken_this_cycle += 1
                
                # Determine if this was a true rest or just a cooldown that was waiting for a slot.
                is_waiting_after_cooldown = profile.get('rest_reason') == "Waiting for group capacity"

                if not self.dry_run:
                    # When un-resting from a long rest, reset counters to give it a fresh start.
                    # Do not reset for COOLDOWN or a profile that was waiting after cooldown.
                    self.manager.update_profile_state(profile['name'], self.manager.STATE_ACTIVE, "Rest period completed")
                    if profile['state'] == self.manager.STATE_RESTING and not is_waiting_after_cooldown:
                        self.manager.reset_profile_counters(profile['name'])
                
                # Update the shared map to reflect the change immediately for this cycle.
                all_profiles_map[profile_name]['state'] = self.manager.STATE_ACTIVE
                if profile['state'] == self.manager.STATE_RESTING and not is_waiting_after_cooldown:
                    all_profiles_map[profile_name]['success_count'] = 0
                    all_profiles_map[profile_name]['failure_count'] = 0
                    all_profiles_map[profile_name]['tolerated_error_count'] = 0
                    all_profiles_map[profile_name]['download_count'] = 0
                    all_profiles_map[profile_name]['download_error_count'] = 0

    def enforce_failure_rate_policy(self, profile, max_failure_rate, min_requests):
        if max_failure_rate <= 0:
            return
            
        success = profile.get('global_success_count', 0)
        failure = profile.get('global_failure_count', 0)
        total = success + failure
        
        if total < min_requests:
            return
        
        current_failure_rate = failure / total if total > 0 else 0
        
        if current_failure_rate >= max_failure_rate:
            reason = f"Global failure rate {current_failure_rate:.2f} >= threshold {max_failure_rate} ({int(failure)}/{int(total)} failures)"
            logger.warning(f"Banning profile '{profile['name']}' due to high failure rate: {reason}")
            if not self.dry_run:
                self.manager.update_profile_state(profile['name'], self.manager.STATE_BANNED, reason)
            self.actions_taken_this_cycle += 1

    def enforce_rest_policy(self, profile, rest_after_requests, rest_duration_minutes):
        if not rest_after_requests or rest_after_requests <= 0 or not rest_duration_minutes or rest_duration_minutes <= 0:
            return

        total_requests = (
            int(profile.get('success_count', 0)) + 
            int(profile.get('failure_count', 0)) + 
            int(profile.get('tolerated_error_count', 0)) +
            int(profile.get('download_count', 0)) +
            int(profile.get('download_error_count', 0))
        )
        
        if total_requests >= rest_after_requests:
            reason = f"Request count {total_requests} >= threshold {rest_after_requests}"
            logger.info(f"Resting profile '{profile['name']}' for {rest_duration_minutes}m: {reason}")
            self.actions_taken_this_cycle += 1
            if not self.dry_run:
                self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, reason)
                rest_until = time.time() + rest_duration_minutes * 60
                self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until))

    def enforce_stale_lock_cleanup(self, max_lock_seconds):
        """Finds and unlocks profiles with stale locks."""
        if self.dry_run:
            logger.info(f"[Dry Run] Would check for and clean up locks older than {max_lock_seconds} seconds.")
            return
        
        cleaned_count = self.manager.cleanup_stale_locks(max_lock_seconds)
        if cleaned_count > 0:
            self.actions_taken_this_cycle += cleaned_count

    def enforce_profile_group_policies(self, profile_groups, all_profiles_map):
        """
        Manages profiles within defined groups. This includes:
        1. Rotating out profiles that have met their request limit.
        2. Healing the group by ensuring no more than `max_active_profiles` are active.
        3. Initializing the group by activating a profile if none are active.
        
        This method operates on and modifies the `all_profiles_map` passed to it.
        """
        if not profile_groups:
            return

        all_profiles_list = list(all_profiles_map.values())

        for group in profile_groups:
            group_name = group.get('name')
            if not group_name:
                logger.warning("Found a profile group without a 'name'. Skipping.")
                continue

            profiles_in_group = set()
            if 'profiles' in group:
                profiles_in_group = set(group['profiles'])
            elif 'prefix' in group:
                prefix = group['prefix']
                profiles_in_group = {p['name'] for p in all_profiles_list if p['name'].startswith(prefix)}
            
            if not profiles_in_group:
                logger.warning(f"Profile group '{group_name}' has no matching profiles. Skipping.")
                continue

            # --- Persist group policy to Redis for observability ---
            rotate_after_requests = group.get('rotate_after_requests')
            max_active_profiles = group.get('max_active_profiles')
            if not self.dry_run:
                # This is a non-critical update, so we don't need to check for existence.
                # We just update it on every cycle to ensure it's fresh.
                self.manager.set_profile_group_state(group_name, {
                    'rotate_after_requests': rotate_after_requests,
                    'max_active_profiles': max_active_profiles,
                    'prefix': group.get('prefix') # Store prefix for observability
                })

            # --- 1. Handle Rotation for Active Profiles ---
            rotate_after_requests = group.get('rotate_after_requests')
            if rotate_after_requests and rotate_after_requests > 0:
                # Consider ACTIVE, LOCKED, and COOLDOWN profiles for rotation eligibility.
                eligible_for_rotation_check = [
                    p for p in all_profiles_list
                    if p['name'] in profiles_in_group and p['state'] in [self.manager.STATE_ACTIVE, self.manager.STATE_LOCKED, self.manager.STATE_COOLDOWN]
                ]

                for profile in eligible_for_rotation_check:
                    total_requests = (
                        int(profile.get('success_count', 0)) + 
                        int(profile.get('failure_count', 0)) + 
                        int(profile.get('tolerated_error_count', 0)) +
                        int(profile.get('download_count', 0)) +
                        int(profile.get('download_error_count', 0))
                    )
                    if total_requests >= rotate_after_requests:
                        # If a profile is LOCKED, we can't rotate it yet.
                        # Instead, we update its reason to show that a rotation is pending.
                        if profile['state'] == self.manager.STATE_LOCKED:
                            pending_reason = f"Pending Rotation (requests: {total_requests}/{rotate_after_requests})"
                            # Only update if the reason is not already set, to avoid spamming Redis.
                            if profile.get('reason') != pending_reason:
                                logger.info(f"Profile '{profile['name']}' in group '{group_name}' is due for rotation but is LOCKED. Marking as pending.")
                                self.actions_taken_this_cycle += 1
                                if not self.dry_run:
                                    self.manager.update_profile_field(profile['name'], 'reason', pending_reason)
                            else:
                                logger.debug(f"Profile '{profile['name']}' in group '{group_name}' is due for rotation but is currently LOCKED. Already marked as pending.")
                            continue
                        
                        # If the profile is ACTIVE or in COOLDOWN, we can rotate it immediately.
                        reason = f"Rotated after {total_requests} requests (limit: {rotate_after_requests})"
                        logger.info(f"Rotating profile '{profile['name']}' in group '{group_name}': {reason}")
                        self.actions_taken_this_cycle += 1

                        wait_for_downloads = group.get('wait_download_finish_per_profile', False)
                        
                        new_reason = reason
                        rest_until_ts = 0

                        if wait_for_downloads:
                            new_reason = "waiting_downloads"
                            logger.info(f"Profile '{profile['name']}' will wait for pending downloads to complete.")
                        else:
                            rest_duration_minutes = group.get('rest_duration_minutes_on_rotation')
                            if rest_duration_minutes and rest_duration_minutes > 0:
                                rest_until_ts = time.time() + rest_duration_minutes * 60

                        if not self.dry_run:
                            self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, new_reason)
                            
                            if wait_for_downloads:
                                self.manager.update_profile_field(profile['name'], 'wait_started_at', str(time.time()))
                                # Set rest_until to 0 to indicate it's not a time-based rest
                                self.manager.update_profile_field(profile['name'], 'rest_until', '0')
                            elif rest_until_ts > 0:
                                self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))

                            # Reset all session counters for the next active cycle
                            self.manager.reset_profile_counters(profile['name'])

                        # Update our local map so subsequent policies in this cycle see the change immediately.
                        all_profiles_map[profile['name']]['state'] = self.manager.STATE_RESTING
                        all_profiles_map[profile['name']]['rest_reason'] = new_reason
                        if wait_for_downloads:
                            all_profiles_map[profile['name']]['wait_started_at'] = time.time()
                            all_profiles_map[profile['name']]['rest_until'] = 0
                        else:
                            all_profiles_map[profile['name']]['rest_until'] = rest_until_ts
                        all_profiles_map[profile['name']]['success_count'] = 0
                        all_profiles_map[profile['name']]['failure_count'] = 0
                        all_profiles_map[profile['name']]['tolerated_error_count'] = 0
                        all_profiles_map[profile['name']]['download_count'] = 0
                        all_profiles_map[profile['name']]['download_error_count'] = 0

            # --- 2. Self-Healing: Enforce max_active_profiles ---
            max_active = group.get('max_active_profiles', 1)
            
            # Get the current list of active/locked profiles from our potentially modified local map
            # A profile is considered "active" for group limits if it is ACTIVE, LOCKED, or in COOLDOWN.
            current_active_or_locked_profiles = [
                p for name, p in all_profiles_map.items() 
                if name in profiles_in_group and p['state'] in [self.manager.STATE_ACTIVE, self.manager.STATE_LOCKED, self.manager.STATE_COOLDOWN]
            ]
            
            num_active_or_locked = len(current_active_or_locked_profiles)
            if num_active_or_locked > max_active:
                logger.warning(f"Healing group '{group_name}': Found {num_active_or_locked} active/locked profiles, but max is {max_active}. Resting excess ACTIVE profiles.")
                
                # We can only rest profiles that are in the ACTIVE state, not LOCKED.
                profiles_that_can_be_rested = [p for p in current_active_or_locked_profiles if p['state'] == self.manager.STATE_ACTIVE]

                # Sort to determine which profiles to rest. We prefer to rest profiles
                # that have been used more. As a tie-breaker (especially for profiles
                # with 0 requests), we rest the one that has been active the longest
                # (oldest last_used timestamp).
                profiles_that_can_be_rested.sort(key=lambda p: p.get('last_used', 0)) # Oldest first
                profiles_that_can_be_rested.sort(key=lambda p: (
                    p.get('success_count', 0) + p.get('failure_count', 0) +
                    p.get('tolerated_error_count', 0) +
                    p.get('download_count', 0) + p.get('download_error_count', 0)
                ), reverse=True) # Most requests first
                
                num_to_rest = num_active_or_locked - max_active
                profiles_to_rest = profiles_that_can_be_rested[:num_to_rest]
                for profile in profiles_to_rest:
                    req_count = (
                        profile.get('success_count', 0) + profile.get('failure_count', 0) +
                        profile.get('tolerated_error_count', 0) +
                        profile.get('download_count', 0) +
                        profile.get('download_error_count', 0)
                    )
                    logger.warning(f"Healing group '{group_name}': Resting profile '{profile['name']}' (request count: {req_count}).")
                    self.actions_taken_this_cycle += 1
                    if not self.dry_run:
                        self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, "Group max_active healing")
                        # Give it a rest time of 0, so it's immediately eligible for activation
                        # by the unrest logic if the group has capacity.
                        rest_until_ts = 0
                        self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
            
            # --- 3. Initialization: Activate profiles if below capacity ---
            # This is a fallback for initialization or if all profiles were rested/banned.
            # The primary activation mechanism is in `enforce_unrest_policy`.
            elif num_active_or_locked < max_active:
                # Check if there are any non-active, non-banned, non-locked profiles to activate.
                eligible_profiles = [p for name, p in all_profiles_map.items() if name in profiles_in_group and p['state'] not in [self.manager.STATE_ACTIVE, self.manager.STATE_BANNED, self.manager.STATE_LOCKED]]
                if eligible_profiles:
                    # This is a simple initialization case. We don't activate here because
                    # `enforce_unrest_policy` will handle it more intelligently based on rest times.
                    # This block ensures that on the very first run, a group doesn't sit empty.
                    if num_active_or_locked == 0:
                         logger.debug(f"Group '{group_name}' has no active profiles. `enforce_unrest_policy` will attempt to activate one.")

    def enforce_proxy_group_rotation(self, proxy_groups):
        """Manages mutually exclusive work cycles for proxies within defined groups."""
        if not proxy_groups:
            return

        group_names = [g['name'] for g in proxy_groups if g.get('name')]
        if not group_names:
            return
            
        group_states = self.manager.get_proxy_group_states(group_names)
        now = time.time()

        for group in proxy_groups:
            group_name = group.get('name')
            if not group_name:
                logger.warning("Found a proxy group without a 'name'. Skipping.")
                continue

            proxies_in_group = group.get('proxies', [])
            if not proxies_in_group:
                logger.warning(f"Proxy group '{group_name}' has no proxies defined. Skipping.")
                continue

            work_minutes = group.get('work_minutes_per_proxy')
            if not work_minutes or work_minutes <= 0:
                logger.warning(f"Proxy group '{group_name}' is missing 'work_minutes_per_proxy'. Skipping.")
                continue
            
            if not self.dry_run:
                for proxy_url in proxies_in_group:
                    self.manager.set_proxy_group_membership(proxy_url, group_name, work_minutes)

            work_duration_seconds = work_minutes * 60
            state = group_states.get(group_name, {})
            
            if not state:
                # First run for this group, initialize it
                logger.info(f"Initializing new proxy group '{group_name}'. Activating first proxy '{proxies_in_group[0]}'.")
                self.actions_taken_this_cycle += 1
                active_proxy_index = 0
                next_rotation_ts = now + work_duration_seconds
                
                if not self.dry_run:
                    # Activate the first, rest the others
                    self.manager.set_proxy_state(proxies_in_group[0], self.manager.STATE_ACTIVE)
                    for i, proxy_url in enumerate(proxies_in_group):
                        if i != active_proxy_index:
                            # Rest indefinitely; group logic will activate it when its turn comes.
                            self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_duration_minutes=99999)
                    
                    self.manager.set_proxy_group_state(group_name, active_proxy_index, next_rotation_ts)
            
            elif now >= state.get('next_rotation_timestamp', 0):
                # Time to rotate
                current_active_index = state.get('active_proxy_index', 0)
                next_active_index = (current_active_index + 1) % len(proxies_in_group)
                
                old_active_proxy = proxies_in_group[current_active_index]
                new_active_proxy = proxies_in_group[next_active_index]
                
                logger.info(f"Rotating proxy group '{group_name}': Deactivating '{old_active_proxy}', Activating '{new_active_proxy}'.")
                self.actions_taken_this_cycle += 1

                next_rotation_ts = now + work_duration_seconds

                if not self.dry_run:
                    # Rest the old proxy
                    self.manager.set_proxy_state(old_active_proxy, self.manager.STATE_RESTING, rest_duration_minutes=99999)
                    # Activate the new one
                    self.manager.set_proxy_state(new_active_proxy, self.manager.STATE_ACTIVE)
                    # Update group state
                    self.manager.set_proxy_group_state(group_name, next_active_index, next_rotation_ts)

    def enforce_proxy_work_rest_cycle(self, args):
        """Enforces a work/rest cycle on proxies based on time."""
        work_minutes = args.proxy_work_minutes
        rest_minutes = args.proxy_rest_duration_minutes

        if not work_minutes or work_minutes <= 0 or not rest_minutes or rest_minutes <= 0:
            return
        
        # Get a flat list of all proxies managed by groups, so we can ignore them.
        proxy_groups = getattr(args, 'proxy_groups', [])
        grouped_proxies = set()
        if proxy_groups:
            for group in proxy_groups:
                for proxy_url in group.get('proxies', []):
                    grouped_proxies.add(proxy_url)

        all_profiles = self.manager.list_profiles()
        if not all_profiles:
            return
            
        unique_proxies = sorted(list(set(p['proxy'] for p in all_profiles if p.get('proxy'))))
        
        # Filter out proxies that are managed by the group rotation logic
        proxies_to_manage = [p for p in unique_proxies if p not in grouped_proxies]
        if not proxies_to_manage:
            logger.debug("All unique proxies are managed by proxy groups. Skipping individual work/rest cycle enforcement.")
            return

        proxy_states = self.manager.get_proxy_states(proxies_to_manage)
        now = time.time()
        
        for proxy_url, state_data in proxy_states.items():
            state = state_data.get('state', self.manager.STATE_ACTIVE)
            
            # Un-rest logic
            if state == self.manager.STATE_RESTING:
                rest_until = state_data.get('rest_until', 0)
                if now >= rest_until:
                    logger.info(f"Activating proxy '{proxy_url}' (rest period complete).")
                    self.actions_taken_this_cycle += 1
                    if not self.dry_run:
                        self.manager.set_proxy_state(proxy_url, self.manager.STATE_ACTIVE)

                    # Also activate any profiles that were resting due to this proxy
                    profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
                    for profile in profiles_for_proxy:
                        if profile['state'] == self.manager.STATE_RESTING and profile.get('rest_reason') == self.PROXY_REST_REASON:
                            logger.info(f"Activating profile '{profile['name']}' as its proxy '{proxy_url}' is now active.")
                            self.actions_taken_this_cycle += 1
                            if not self.dry_run:
                                self.manager.update_profile_state(profile['name'], self.manager.STATE_ACTIVE, "Proxy activated")
                else:
                    # Proxy is still resting. Ensure any of its profiles that are ACTIVE are moved to RESTING.
                    # This catches profiles that were unlocked while their proxy was resting.
                    rest_until_ts = state_data.get('rest_until', 0)
                    profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
                    for profile in profiles_for_proxy:
                        if profile['state'] == self.manager.STATE_ACTIVE:
                            logger.info(f"Resting profile '{profile['name']}' as its proxy '{proxy_url}' is resting.")
                            self.actions_taken_this_cycle += 1
                            if not self.dry_run:
                                self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, self.PROXY_REST_REASON)
                                self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
            
            # Rest logic
            elif state == self.manager.STATE_ACTIVE:
                work_start = state_data.get('work_start_timestamp', 0)
                if work_start == 0: # Proxy was just created, start its work cycle
                    if not self.dry_run:
                        self.manager.set_proxy_state(proxy_url, self.manager.STATE_ACTIVE)
                    continue

                work_duration_seconds = work_minutes * 60
                active_duration = now - work_start
                logger.debug(f"Proxy '{proxy_url}' has been active for {active_duration:.0f}s (limit: {work_duration_seconds}s).")
                if active_duration >= work_duration_seconds:
                    logger.info(f"Resting proxy '{proxy_url}' for {rest_minutes}m (work period of {work_minutes}m complete).")
                    self.actions_taken_this_cycle += 1
                    
                    rest_until_ts = time.time() + rest_minutes * 60
                    if not self.dry_run:
                        self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_minutes)
                    
                    # Also rest any active profiles using this proxy
                    profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
                    for profile in profiles_for_proxy:
                        if profile['state'] == self.manager.STATE_ACTIVE:
                            logger.info(f"Resting profile '{profile['name']}' as its proxy '{proxy_url}' is resting.")
                            self.actions_taken_this_cycle += 1
                            if not self.dry_run:
                                self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, self.PROXY_REST_REASON)
                                self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))

    def enforce_max_proxy_active_time(self, args):
        """
        Enforces a global maximum active time for any proxy, regardless of group membership.
        This acts as a safety net to prevent a proxy from being stuck in an ACTIVE state.
        """
        max_active_minutes = args.max_global_proxy_active_minutes
        rest_minutes = args.rest_duration_on_max_active

        if not max_active_minutes or max_active_minutes <= 0:
            return

        all_profiles = self.manager.list_profiles()
        if not all_profiles:
            return
            
        unique_proxies = sorted(list(set(p['proxy'] for p in all_profiles if p.get('proxy'))))
        if not unique_proxies:
            return

        proxy_states = self.manager.get_proxy_states(unique_proxies)
        now = time.time()
        
        for proxy_url, state_data in proxy_states.items():
            if state_data.get('state') == self.manager.STATE_ACTIVE:
                work_start = state_data.get('work_start_timestamp', 0)
                if work_start == 0:
                    continue # Just activated, timestamp not set yet.

                active_duration_seconds = now - work_start
                max_active_seconds = max_active_minutes * 60

                if active_duration_seconds >= max_active_seconds:
                    reason = f"Exceeded max active time of {max_active_minutes}m"
                    logger.warning(f"Resting proxy '{proxy_url}' for {rest_minutes}m: {reason}")
                    self.actions_taken_this_cycle += 1
                    
                    rest_until_ts = now + rest_minutes * 60
                    if not self.dry_run:
                        self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_minutes)
                    
                    # Also rest any active profiles using this proxy
                    profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
                    for profile in profiles_for_proxy:
                        if profile['state'] == self.manager.STATE_ACTIVE:
                            logger.info(f"Resting profile '{profile['name']}' as its proxy '{proxy_url}' is resting due to max active time.")
                            self.actions_taken_this_cycle += 1
                            if not self.dry_run:
                                self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, self.PROXY_REST_REASON)
                                self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))

    def enforce_proxy_policies(self, args):
        proxy_ban_enabled = args.proxy_ban_on_failures and args.proxy_ban_on_failures > 0
        proxy_rate_limit_enabled = getattr(args, 'proxy_rate_limit_requests', 0) > 0
        if not proxy_ban_enabled and not proxy_rate_limit_enabled:
            return

        all_profiles = self.manager.list_profiles()
        if not all_profiles:
            return
            
        unique_proxies = sorted(list(set(p['proxy'] for p in all_profiles if p.get('proxy'))))

        if not unique_proxies:
            return
            
        logger.debug(f"Checking proxy policies for {len(unique_proxies)} unique proxies...")
        
        for proxy_url in unique_proxies:
            profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
            if self.enforce_proxy_failure_burst_policy(
                proxy_url,
                profiles_for_proxy,
                args.proxy_ban_on_failures, 
                args.proxy_ban_window_minutes
            ):
                continue # Banned, no need for other checks
            
            self.enforce_proxy_rate_limit_policy(
                proxy_url,
                profiles_for_proxy,
                getattr(args, 'proxy_rate_limit_requests', 0),
                getattr(args, 'proxy_rate_limit_window_minutes', 0),
                getattr(args, 'proxy_rate_limit_rest_duration_minutes', 0)
            )
    
    def enforce_proxy_failure_burst_policy(self, proxy_url, profiles_for_proxy, max_failures, window_minutes):
        if not max_failures or not window_minutes or max_failures <= 0 or window_minutes <= 0:
            return False

        window_seconds = window_minutes * 60
        failure_count = self.manager.get_proxy_activity_rate(proxy_url, 'failure', window_seconds)
        
        if failure_count >= max_failures:
            reason = f"Proxy failure burst: {failure_count} failures in last {window_minutes}m (threshold: {max_failures})"
            logger.warning(f"Banning {len(profiles_for_proxy)} profile(s) on proxy '{proxy_url}' due to failure burst: {reason}")
            self.actions_taken_this_cycle += 1
            
            if not self.dry_run:
                for profile in profiles_for_proxy:
                    # Don't re-ban already banned profiles
                    if profile['state'] != self.manager.STATE_BANNED:
                        self.manager.update_profile_state(profile['name'], self.manager.STATE_BANNED, reason)
            return True # Indicates action was taken
        return False

    def enforce_proxy_rate_limit_policy(self, proxy_url, profiles_for_proxy, max_requests, window_minutes, rest_duration_minutes):
        if not max_requests or not window_minutes or max_requests <= 0 or window_minutes <= 0:
            return False

        window_seconds = window_minutes * 60
        # Count all successful activities for the proxy
        activity_count = (
            self.manager.get_proxy_activity_rate(proxy_url, 'success', window_seconds) +
            self.manager.get_proxy_activity_rate(proxy_url, 'download', window_seconds)
        )

        if activity_count >= max_requests:
            reason = f"Proxy rate limit hit: {activity_count} requests in last {window_minutes}m (limit: {max_requests})"
            logger.info(f"Resting proxy '{proxy_url}' for {rest_duration_minutes}m: {reason}")
            self.actions_taken_this_cycle += 1
            
            if not self.dry_run:
                self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_duration_minutes)
            return True # Indicates action was taken
        return False

def add_policy_enforcer_parser(subparsers):
    """Adds the parser for the 'policy-enforcer' command."""
    parser = subparsers.add_parser(
        'policy-enforcer',
        description='Apply policies to profiles (ban, rest, etc.).',
        formatter_class=argparse.RawTextHelpFormatter,
        help='Apply policies to profiles (ban, rest, etc.).'
    )
    
    parser.add_argument('--policy', '--policy-file', dest='policy_file', help='Path to a YAML policy file to load default settings from.')
    parser.add_argument('--env-file', help='Path to a .env file to load environment variables from.')
    parser.add_argument('--redis-host', default=None, help='Redis host. Defaults to REDIS_HOST or MASTER_HOST_IP env var, or localhost.')
    parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Defaults to REDIS_PORT env var, or 6379.')
    parser.add_argument('--redis-password', default=None, help='Redis password. Defaults to REDIS_PASSWORD env var.')
    parser.add_argument('--env', default=None, help="Default environment name for Redis key prefix. Used if --auth-env or --download-env are not specified. Overrides policy file setting.")
    parser.add_argument('--auth-env', help="Override the environment for the Auth simulation.")
    parser.add_argument('--download-env', help="Override the environment for the Download simulation.")
    parser.add_argument('--legacy', action='store_true', help="Use legacy key prefix ('profile_mgmt_') without environment.")
    parser.add_argument('--key-prefix', default=None, help='Explicit key prefix for Redis. Overrides --env, --legacy and any defaults.')
    parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
    parser.add_argument('--dry-run', action='store_true', help='Show what would be done without making changes.')
    
    # Policy arguments
    policy_group = parser.add_argument_group('Policy Rules')
    policy_group.add_argument('--max-failure-rate', type=float, default=None,
                              help='Ban a profile if its failure rate exceeds this value (0.0 to 1.0). Default: 0.5')
    policy_group.add_argument('--min-requests-for-rate', type=int, default=None,
                              help='Minimum total requests before failure rate is calculated. Default: 20')
    policy_group.add_argument('--ban-on-failures', type=int, default=None,
                              help='Ban a profile if it has this many failures within the time window (0 to disable). Default: 0')
    policy_group.add_argument('--ban-on-failures-window-minutes', type=int, default=None,
                              help='The time window in minutes for the failure burst check. Default: 5')
    policy_group.add_argument('--rest-after-requests', type=int, default=None,
                              help='Move a profile to RESTING after this many total requests (0 to disable). Default: 0')
    policy_group.add_argument('--rest-duration-minutes', type=int, default=None,
                              help='How long a profile should rest. Default: 15')
    policy_group.add_argument('--rate-limit-requests', type=int, default=None,
                              help='Rest a profile if it exceeds this many requests in the time window (0 to disable).')
    policy_group.add_argument('--rate-limit-window-minutes', type=int, default=None,
                              help='The time window in minutes for the rate limit check.')
    policy_group.add_argument('--rate-limit-rest-duration-minutes', type=int, default=None,
                              help='How long a profile should rest after hitting the rate limit.')
    policy_group.add_argument('--unlock-stale-locks-after-seconds', type=int, default=None,
                              help='Unlock profiles that have been in a LOCKED state for more than this many seconds (0 to disable). Default: 120')
    
    proxy_policy_group = parser.add_argument_group('Proxy Policy Rules')
    proxy_policy_group.add_argument('--proxy-work-minutes', type=int, default=None,
                                    help='Work duration for a proxy before it rests (0 to disable). Default: 0')
    proxy_policy_group.add_argument('--proxy-rest-duration-minutes', type=int, default=None,
                                    help='Rest duration for a proxy after its work period. Default: 0')
    proxy_policy_group.add_argument('--proxy-ban-on-failures', type=int, default=None,
                                    help='Ban a proxy (and all its profiles) if it has this many failures within the time window (0 to disable). Default: 0')
    proxy_policy_group.add_argument('--proxy-ban-window-minutes', type=int, default=None,
                                    help='The time window in minutes for the proxy failure burst check. Default: 10')
    proxy_policy_group.add_argument('--proxy-rate-limit-requests', type=int, default=None,
                                    help='Rest a proxy if it exceeds this many requests in the time window (0 to disable).')
    proxy_policy_group.add_argument('--proxy-rate-limit-window-minutes', type=int, default=None,
                                    help='The time window in minutes for the proxy rate limit check.')
    proxy_policy_group.add_argument('--proxy-rate-limit-rest-duration-minutes', type=int, default=None,
                                    help='How long a proxy should rest after hitting the rate limit.')
    proxy_policy_group.add_argument('--max-global-proxy-active-minutes', type=int, default=None,
                                    help='Global maximum time a proxy can be active before being rested (0 to disable). Acts as a safety net. Default: 0')
    proxy_policy_group.add_argument('--rest-duration-on-max-active', type=int, default=None,
                                    help='How long a proxy should rest after hitting the global max active time. Default: 10')

    # Execution control
    exec_group = parser.add_argument_group('Execution Control')
    exec_group.add_argument('--live', action='store_true', help='Run continuously, applying policies periodically.')
    exec_group.add_argument('--interval-seconds', type=int, default=None,
                            help='When in --live mode, how often to apply policies. Default: 60')
    exec_group.add_argument('--auth-only', action='store_true', help='Run enforcer for the auth simulation only.')
    exec_group.add_argument('--download-only', action='store_true', help='Run enforcer for the download simulation only.')

    return parser

def sync_cross_simulation(auth_manager, download_manager, sync_config, dry_run=False):
    """Synchronize profile states between auth and download simulations."""
    if not sync_config:
        return

    profile_links = sync_config.get('profile_links', [])
    sync_states = sync_config.get('sync_states', [])
    sync_rotation = sync_config.get('sync_rotation', False)
    enforce_auth_lead = sync_config.get('enforce_auth_lead', False)

    if not profile_links:
        return

    # --- Get all profiles once for efficiency ---
    all_auth_profiles = {p['name']: p for p in auth_manager.list_profiles()}
    all_download_profiles = {p['name']: p for p in download_manager.list_profiles()}

    # --- State and Rotation Sync (handles prefixes correctly) ---
    for link in profile_links:
        auth_prefix = link.get('auth')
        download_prefix = link.get('download')
        if not auth_prefix or not download_prefix:
            continue

        auth_profiles_in_group = [p for name, p in all_auth_profiles.items() if name.startswith(auth_prefix)]

        for auth_profile in auth_profiles_in_group:
            # Assume 1-to-1 name mapping (e.g., auth 'user1_0' maps to download 'user1_0')
            download_profile_name = auth_profile['name']
            download_profile = all_download_profiles.get(download_profile_name)

            if not download_profile:
                logger.debug(f"Auth profile '{auth_profile['name']}' has no corresponding download profile.")
                continue

            auth_state = auth_profile.get('state')
            download_state = download_profile.get('state')

            # Sync states from auth to download
            if enforce_auth_lead and auth_state in sync_states and download_state != auth_state:
                auth_reason = auth_profile.get('reason', '')
                # If auth profile is waiting for downloads, we must NOT sync the RESTING state to the download profile,
                # as that would prevent it from processing the very downloads we are waiting for.
                if auth_state == auth_manager.STATE_RESTING and auth_reason == 'waiting_downloads':
                    logger.debug(f"Auth profile '{auth_profile['name']}' is waiting for downloads. Skipping state sync to download profile to prevent deadlock.")
                else:
                    logger.info(f"Syncing download profile '{download_profile_name}' to state '{auth_state}' (auth lead)")
                    if not dry_run:
                        reason_to_sync = auth_reason or 'Synced from auth'
                        download_manager.update_profile_state(download_profile_name, auth_state, f"Synced from auth: {reason_to_sync}")
                        if auth_state == auth_manager.STATE_RESTING:
                            auth_rest_until = auth_profile.get('rest_until')
                            if auth_rest_until:
                                download_manager.update_profile_field(download_profile_name, 'rest_until', str(auth_rest_until))

            # Handle rotation sync
            if sync_rotation:
                auth_reason = auth_profile.get('rest_reason', '')
                
                # If auth profile is waiting for downloads, we must NOT sync the RESTING state to the download profile,
                # as that would prevent it from processing the very downloads we are waiting for.
                if auth_reason == 'waiting_downloads':
                    logger.debug(f"Auth profile '{auth_profile['name']}' is waiting for downloads. Skipping rotation sync to download profile to prevent deadlock.")
                elif auth_state == auth_manager.STATE_RESTING and 'rotate' in auth_reason.lower():
                    if download_state != download_manager.STATE_RESTING:
                        logger.info(f"Rotating download profile '{download_profile_name}' due to auth rotation")
                        if not dry_run:
                            download_manager.update_profile_state(download_profile_name, download_manager.STATE_RESTING, f"Rotated due to auth rotation: {auth_reason}")
                            auth_rest_until = auth_profile.get('rest_until')
                            if auth_rest_until:
                                download_manager.update_profile_field(download_profile_name, 'rest_until', str(auth_rest_until))

    # --- Active Profile Sync ---
    sync_active = sync_config.get('sync_active_profile', False)
    sync_waiting_downloads = sync_config.get('sync_waiting_downloads', False)
    
    if not (sync_active or sync_waiting_downloads):
        return

    logger.debug("Syncing active profiles from Auth to Download simulation...")

    # Get profiles that should be active in the download simulation
    target_active_download_profiles = set()
    
    # 1. Add profiles that are active in auth simulation (if sync_active is enabled)
    if sync_active:
        active_auth_profiles = [p for p in all_auth_profiles.values() if p['state'] in [auth_manager.STATE_ACTIVE, auth_manager.STATE_LOCKED]]
        for auth_profile in active_auth_profiles:
            target_active_download_profiles.add(auth_profile['name'])
    
    # 2. Add profiles that are waiting for downloads to complete (if sync_waiting_downloads is enabled)
    if sync_waiting_downloads:
        waiting_auth_profiles = [p for p in all_auth_profiles.values() 
                               if p['state'] == auth_manager.STATE_RESTING 
                               and p.get('rest_reason') == 'waiting_downloads']
        for auth_profile in waiting_auth_profiles:
            target_active_download_profiles.add(auth_profile['name'])
            logger.debug(f"Auth profile '{auth_profile['name']}' is waiting for downloads. Ensuring matching download profile is active.")
    
    if not target_active_download_profiles:
        logger.debug("No auth profiles found that need active download profiles.")
        return

    # Get download profile group info from Redis
    dl_group_state_keys = [k for k in download_manager.redis.scan_iter(f"{download_manager.key_prefix}profile_group_state:*")]
    dl_group_names = [k.split(':')[-1] for k in dl_group_state_keys]
    dl_group_states = download_manager.get_profile_group_states(dl_group_names)
    
    dl_profile_to_group = {}
    for name, state in dl_group_states.items():
        prefix = state.get('prefix')
        if prefix:
            for p_name in all_download_profiles:
                if p_name.startswith(prefix):
                    dl_profile_to_group[p_name] = {'name': name, 'max_active': state.get('max_active_profiles', 1)}

    # Activate download profiles that should be active but aren't
    for target_profile_name in target_active_download_profiles:
        download_profile = all_download_profiles.get(target_profile_name)
        if not download_profile:
            logger.warning(f"Auth profile '{target_profile_name}' needs an active download profile, but no corresponding download profile found.")
            continue

        if download_profile['state'] not in [download_manager.STATE_ACTIVE, download_manager.STATE_LOCKED]:
            is_from_cooldown = download_profile['state'] == download_manager.STATE_COOLDOWN
            log_msg_suffix = " (from COOLDOWN)" if is_from_cooldown else ""
            logger.info(f"Syncing active state: Activating download profile '{target_profile_name}' to match auth requirements{log_msg_suffix}.")
            if not dry_run:
                download_manager.update_profile_state(target_profile_name, download_manager.STATE_ACTIVE, "Synced from auth requirements")
                # Only reset counters if it's coming from a long-term rest, not a short cooldown.
                if not is_from_cooldown:
                    download_manager.reset_profile_counters(target_profile_name)

    # Deactivate any download profiles that are active but shouldn't be
    for dl_profile_name, dl_profile in all_download_profiles.items():
        if dl_profile['state'] == download_manager.STATE_ACTIVE and dl_profile_name not in target_active_download_profiles:
            group_info = dl_profile_to_group.get(dl_profile_name)
            if group_info:
                logger.info(f"Syncing active state: Resting download profile '{dl_profile_name}' as it is no longer the active profile in its group.")
                if not dry_run:
                    download_manager.update_profile_state(dl_profile_name, download_manager.STATE_RESTING, "Synced rotation from auth")
                    download_manager.update_profile_field(dl_profile_name, 'rest_until', '0')

def main_policy_enforcer(args):
    """Main dispatcher for 'policy-enforcer' command."""
    policy = {}
    if args.policy_file:
        if not yaml:
            logger.error("Cannot load policy file because PyYAML is not installed.")
            return 1
        try:
            with open(args.policy_file, 'r') as f:
                policy = yaml.safe_load(f) or {}
        except (IOError, yaml.YAMLError) as e:
            logger.error(f"Failed to load or parse policy file {args.policy_file}: {e}")
            return 1

    class Config:
        def __init__(self, cli_args, policy_defaults, code_defaults):
            for key, code_default in code_defaults.items():
                cli_val = getattr(cli_args, key, None)
                policy_val = policy_defaults.get(key)
                if cli_val is not None:
                    setattr(self, key, cli_val)
                elif policy_val is not None:
                    setattr(self, key, policy_val)
                else:
                    setattr(self, key, code_default)

    code_defaults = {
        'max_failure_rate': 0.0, 'min_requests_for_rate': 20, 'ban_on_failures': 0,
        'ban_on_failures_window_minutes': 5, 'rest_after_requests': 0,
        'rest_duration_minutes': 15,
        'rate_limit_requests': 0, 'rate_limit_window_minutes': 60, 'rate_limit_rest_duration_minutes': 5,
        'proxy_work_minutes': 0,
        'proxy_rest_duration_minutes': 0, 'proxy_ban_on_failures': 0,
        'proxy_ban_window_minutes': 10,
        'proxy_rate_limit_requests': 0, 'proxy_rate_limit_window_minutes': 60, 'proxy_rate_limit_rest_duration_minutes': 10,
        'unlock_stale_locks_after_seconds': 120,
        'unlock_cooldown_seconds': 0,
        'max_global_proxy_active_minutes': 0, 'rest_duration_on_max_active': 10,
        'interval_seconds': 60, 'proxy_groups': [], 'profile_groups': []
    }

    sim_params = policy.get('simulation_parameters', {})
    env_file_from_policy = sim_params.get('env_file')

    if load_dotenv:
        env_file = args.env_file or env_file_from_policy
        if not env_file and args.env and '.env' in args.env and os.path.exists(args.env):
            print(f"WARNING: --env should be an environment name, not a file path. Treating '{args.env}' as --env-file.", file=sys.stderr)
            env_file = args.env
        if env_file and load_dotenv(env_file):
            print(f"Loaded environment variables from {env_file}", file=sys.stderr)
        elif args.env_file and not os.path.exists(args.env_file):
            print(f"ERROR: The specified --env-file was not found: {args.env_file}", file=sys.stderr)
            return 1

    redis_host = args.redis_host or os.getenv('REDIS_HOST', os.getenv('MASTER_HOST_IP', 'localhost'))
    redis_port = args.redis_port if args.redis_port is not None else int(os.getenv('REDIS_PORT', 6379))
    redis_password = args.redis_password or os.getenv('REDIS_PASSWORD')

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    signal.signal(signal.SIGINT, handle_shutdown)
    signal.signal(signal.SIGTERM, handle_shutdown)

    enforcer_setups = []

    def setup_enforcer(sim_type, env_cli_arg, policy_config_key, env_policy_key):
        policy_config = policy.get(policy_config_key)
        # Fallback for single-enforcer policy files
        if policy_config is None and sim_type == 'Auth':
            policy_config = policy.get('policy_enforcer_config', {})

        if policy_config is None:
            logger.debug(f"No config block found for {sim_type} simulation ('{policy_config_key}'). Skipping.")
            return None
        
        logger.info(f"Setting up enforcer for {sim_type} simulation...")
        config = Config(args, policy_config, code_defaults)
        
        # Determine the effective environment name with correct precedence:
        # 1. Specific CLI arg (e.g., --auth-env)
        # 2. General CLI arg (--env)
        # 3. Specific policy setting (e.g., simulation_parameters.auth_env)
        # 4. General policy setting (simulation_parameters.env)
        # 5. Hardcoded default ('dev')
        policy_env = sim_params.get(env_policy_key)
        default_policy_env = sim_params.get('env')
        effective_env = env_cli_arg or args.env or policy_env or default_policy_env or 'dev'
        
        logger.info(f"Using environment '{effective_env}' for {sim_type}.")

        if args.key_prefix: key_prefix = args.key_prefix
        elif args.legacy: key_prefix = 'profile_mgmt_'
        else: key_prefix = f"{effective_env}_profile_mgmt_"
        
        manager = ProfileManager(redis_host, redis_port, redis_password, key_prefix)
        enforcer = PolicyEnforcer(manager, dry_run=args.dry_run)
        
        # Write any relevant config to Redis for workers to use
        cooldown = getattr(config, 'unlock_cooldown_seconds', None)
        if cooldown is not None and not args.dry_run:
            # If it's a list or int, convert to JSON string to store in Redis
            manager.set_config('unlock_cooldown_seconds', json.dumps(cooldown))
        
        proxy_work_minutes = getattr(config, 'proxy_work_minutes', None)
        if proxy_work_minutes is not None and not args.dry_run:
            manager.set_config('proxy_work_minutes', proxy_work_minutes)
        
        proxy_rest_duration_minutes = getattr(config, 'proxy_rest_duration_minutes', None)
        if proxy_rest_duration_minutes is not None and not args.dry_run:
            manager.set_config('proxy_rest_duration_minutes', proxy_rest_duration_minutes)

        return {'name': sim_type, 'enforcer': enforcer, 'config': config}

    if not args.download_only:
        auth_setup = setup_enforcer('Auth', args.auth_env, 'auth_policy_enforcer_config', 'auth_env')
        if auth_setup: enforcer_setups.append(auth_setup)
        
    if not args.auth_only:
        download_setup = setup_enforcer('Download', args.download_env, 'download_policy_enforcer_config', 'download_env')
        if download_setup: enforcer_setups.append(download_setup)
    
    if not enforcer_setups:
        logger.error("No policies to enforce. Check policy file and --auth-only/--download-only flags.")
        return 1
        
    # Determine interval. Precedence: CLI -> simulation_parameters -> per-setup config -> code default.
    # The CLI arg is already handled by the Config objects, so we just need to check sim_params.
    sim_params_interval = sim_params.get('interval_seconds')
    if args.interval_seconds is None and sim_params_interval is not None:
        interval = sim_params_interval
    else:
        interval = min(s['config'].interval_seconds for s in enforcer_setups)

    # Get cross-simulation sync configuration
    cross_sync_config = policy.get('cross_simulation_sync', {})
    
    if not args.live:
        for setup in enforcer_setups:
            logger.info(f"--- Applying policies for {setup['name']} Simulation ---")
            setup['enforcer'].apply_policies(setup['config'])
        
        # Apply cross-simulation sync after all policies have been applied
        if cross_sync_config and len(enforcer_setups) == 2:
            # We need to identify which setup is auth and which is download
            # Based on their names
            auth_setup = next((s for s in enforcer_setups if s['name'] == 'Auth'), None)
            download_setup = next((s for s in enforcer_setups if s['name'] == 'Download'), None)
            if auth_setup and download_setup:
                sync_cross_simulation(
                    auth_setup['enforcer'].manager,
                    download_setup['enforcer'].manager,
                    cross_sync_config,
                    dry_run=args.dry_run
                )
        return 0
    
    logger.info(f"Running in live mode. Applying policies every {interval} seconds. Press Ctrl+C to stop.")
    if not args.verbose:
        print("Each '.' represents a check cycle with no actions taken.", file=sys.stderr)

    while not shutdown_event:
        had_action_in_cycle = False
        for setup in enforcer_setups:
            logger.debug(f"--- Applying policies for {setup['name']} Simulation ({setup['enforcer'].manager.key_prefix}) ---")
            if setup['enforcer'].apply_policies(setup['config']):
                had_action_in_cycle = True
        
        # Apply cross-simulation sync after all policies have been applied in this cycle
        if cross_sync_config and len(enforcer_setups) == 2:
            auth_setup = next((s for s in enforcer_setups if s['name'] == 'Auth'), None)
            download_setup = next((s for s in enforcer_setups if s['name'] == 'Download'), None)
            if auth_setup and download_setup:
                sync_cross_simulation(
                    auth_setup['enforcer'].manager,
                    download_setup['enforcer'].manager,
                    cross_sync_config,
                    dry_run=args.dry_run
                )
                # Note: sync_cross_simulation may take actions, but we don't track them for the dot indicator
                # This is fine for now
        
        if had_action_in_cycle:
            if not args.verbose:
                # Print a newline to separate the action logs from subsequent dots
                print(file=sys.stderr)
        else:
            if not args.verbose:
                print(".", end="", file=sys.stderr)
                sys.stderr.flush()

        sleep_end_time = time.time() + interval
        while time.time() < sleep_end_time and not shutdown_event:
            time.sleep(1)
    
    logger.info("Policy enforcer stopped.")
    return 0
