"""
DAG to manage the state of proxies and accounts used by the ytdlp-ops-server.
"""
from __future__ import annotations

import logging
from datetime import datetime
import socket

from airflow.exceptions import AirflowException
from airflow.models.dag import DAG
from airflow.models.dagbag import DagBag
from airflow.models.dagrun import DagRun
from airflow.models.param import Param
from airflow.models.taskinstance import TaskInstance
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago
from airflow.models.variable import Variable
from airflow.providers.redis.hooks.redis import RedisHook
from airflow.utils.session import create_session

# Configure logging
logger = logging.getLogger(__name__)

# Import and apply Thrift exceptions patch for Airflow compatibility
try:
    from thrift_exceptions_patch import patch_thrift_exceptions
    patch_thrift_exceptions()
    logger.info("Applied Thrift exceptions patch for Airflow compatibility.")
except ImportError:
    logger.warning("Could not import thrift_exceptions_patch. Compatibility may be affected.")
except Exception as e:
    logger.error(f"Error applying Thrift exceptions patch: {e}")

# Thrift imports
try:
    from pangramia.yt.exceptions.ttypes import PBServiceException, PBUserException
    from yt_ops_services.client_utils import get_thrift_client, format_timestamp
except ImportError as e:
    logger.critical(f"Could not import Thrift modules: {e}. Ensure yt_ops_services package is installed correctly.")
    # Fail DAG parsing if thrift modules are not available
    raise

DEFAULT_MANAGEMENT_SERVICE_IP = Variable.get("MANAGEMENT_SERVICE_HOST", default_var="envoy-thrift-lb")
DEFAULT_MANAGEMENT_SERVICE_PORT = Variable.get("MANAGEMENT_SERVICE_PORT", default_var=9080)
DEFAULT_REDIS_CONN_ID = "redis_default"


# Helper function to connect to Redis, similar to other DAGs
def _get_redis_client(redis_conn_id: str):
    """Gets a Redis client from an Airflow connection."""
    try:
        # Use the imported RedisHook
        redis_hook = RedisHook(redis_conn_id=redis_conn_id)
        # get_conn returns a redis.Redis client
        return redis_hook.get_conn()
    except Exception as e:
        logger.error(f"Failed to connect to Redis using connection '{redis_conn_id}': {e}")
        # Use the imported AirflowException
        raise AirflowException(f"Redis connection failed: {e}")



def _list_proxy_statuses(client, server_identity):
    """Lists the status of proxies."""
    logger.info(f"Listing proxy statuses for server: {server_identity or 'ALL'}")
    statuses = client.getProxyStatus(server_identity)
    if not statuses:
        logger.info("No proxy statuses found.")
        return

    from tabulate import tabulate
    status_list = []
    # This is forward-compatible: it checks for new attributes before using them.
    has_extended_info = hasattr(statuses[0], 'recentAccounts') or hasattr(statuses[0], 'recentMachines')

    headers = ["Server", "Proxy URL", "Status", "Success", "Failures", "Last Success", "Last Failure"]
    if has_extended_info:
        headers.extend(["Recent Accounts", "Recent Machines"])

    for s in statuses:
        status_item = {
            "Server": s.serverIdentity,
            "Proxy URL": s.proxyUrl,
            "Status": s.status,
            "Success": s.successCount,
            "Failures": s.failureCount,
            "Last Success": format_timestamp(s.lastSuccessTimestamp),
            "Last Failure": format_timestamp(s.lastFailureTimestamp),
        }
        if has_extended_info:
            recent_accounts = getattr(s, 'recentAccounts', [])
            recent_machines = getattr(s, 'recentMachines', [])
            status_item["Recent Accounts"] = "\n".join(recent_accounts) if recent_accounts else "N/A"
            status_item["Recent Machines"] = "\n".join(recent_machines) if recent_machines else "N/A"
        status_list.append(status_item)

    print("\n--- Proxy Statuses ---")
    # The f-string with a newline ensures the table starts on a new line in the logs.
    print(f"\n{tabulate(status_list, headers='keys', tablefmt='grid')}")
    print("----------------------\n")
    if not has_extended_info:
        logger.warning("Server does not seem to support 'recentAccounts' or 'recentMachines' fields yet.")
        print("NOTE: To see Recent Accounts/Machines, the server's `getProxyStatus` method must be updated to return these fields.")


def _list_account_statuses(client, account_id, redis_conn_id):
    """Lists the status of accounts, enriching with live data from Redis."""
    logger.info(f"Listing account statuses for account: {account_id or 'ALL'}")

    redis_client = None
    try:
        redis_client = _get_redis_client(redis_conn_id)
        logger.info("Successfully connected to Redis to fetch detailed account status.")
    except Exception as e:
        logger.warning(f"Could not connect to Redis to get detailed status. Will show basic status. Error: {e}")
        redis_client = None

    try:
        # The thrift method takes accountId (specific) or accountPrefix.
        # If account_id is provided, we use it. If not, we get all by leaving both params as None.
        statuses = client.getAccountStatus(accountId=account_id, accountPrefix=None)
        if not statuses:
            print("\n--- Account Statuses ---\nNo account statuses found.\n------------------------\n")
            return

        from tabulate import tabulate
        status_list = []
        
        for s in statuses:
            status_str = s.status
            # If an account is resting, get the live countdown from Redis for accuracy.
            if redis_client and 'RESTING' in status_str:
                try:
                    status_key = f"account_status:{s.accountId}"
                    # The server stores resting expiry time in 'resting_until'.
                    expiry_ts_bytes = redis_client.hget(status_key, "resting_until")
                    if expiry_ts_bytes:
                        expiry_ts = float(expiry_ts_bytes)
                        now = datetime.now().timestamp()
                        if now >= expiry_ts:
                            status_str = "ACTIVE (was RESTING)"
                        else:
                            remaining_seconds = int(expiry_ts - now)
                            if remaining_seconds > 3600:
                                status_str = f"RESTING (active in {remaining_seconds // 3600}h {remaining_seconds % 3600 // 60}m)"
                            elif remaining_seconds > 60:
                                status_str = f"RESTING (active in {remaining_seconds // 60}m {remaining_seconds % 60}s)"
                            else:
                                status_str = f"RESTING (active in {remaining_seconds}s)"
                except Exception as e:
                    logger.warning(f"Could not parse resting time for {s.accountId} from Redis: {e}. Using server status.")

            # Determine the last activity timestamp for sorting
            last_success = float(s.lastSuccessTimestamp) if s.lastSuccessTimestamp else 0
            last_failure = float(s.lastFailureTimestamp) if s.lastFailureTimestamp else 0
            last_activity = max(last_success, last_failure)

            status_item = {
                "Account ID": s.accountId,
                "Status": status_str,
                "Success": s.successCount,
                "Failures": s.failureCount,
                "Last Success": format_timestamp(s.lastSuccessTimestamp),
                "Last Failure": format_timestamp(s.lastFailureTimestamp),
                "Last Proxy": s.lastUsedProxy or "N/A",
                "Last Machine": s.lastUsedMachine or "N/A",
                "_last_activity": last_activity,  # Add a temporary key for sorting
            }
            status_list.append(status_item)

        # Sort the list by the last activity timestamp in descending order
        status_list.sort(key=lambda item: item.get('_last_activity', 0), reverse=True)

        # Remove the temporary sort key before printing
        for item in status_list:
            del item['_last_activity']

        print("\n--- Account Statuses ---")
        # The f-string with a newline ensures the table starts on a new line in the logs.
        print(f"\n{tabulate(status_list, headers='keys', tablefmt='grid')}")
        print("------------------------\n")
    except (PBServiceException, PBUserException) as e:
        logger.error(f"Failed to get account statuses: {e.message}", exc_info=True)
        print(f"\nERROR: Could not retrieve account statuses. Server returned: {e.message}\n")
    except Exception as e:
        logger.error(f"An unexpected error occurred while getting account statuses: {e}", exc_info=True)
        print(f"\nERROR: An unexpected error occurred: {e}\n")


def manage_system_callable(**context):
    """Main callable to interact with the system management endpoints."""
    params = context["params"]
    entity = params["entity"]
    action = params["action"]

    # For Thrift actions, use the new management host/port
    if entity not in ["airflow_meta"]:
        host = params["management_host"]
        port = params["management_port"]
    else:
        host, port = None, None  # Not needed for meta actions

    server_identity = params.get("server_identity")
    proxy_url = params.get("proxy_url")
    account_id = params.get("account_id")

    # --- Validate Action/Entity Combination and Parameters ---
    valid_actions = {
        "proxy": ["list_with_status", "ban", "unban", "ban_all", "unban_all", "delete_from_redis"],
        "account": ["list_with_status", "ban", "unban", "unban_all", "delete_from_redis"],
        "accounts_and_proxies": ["list_with_status", "ban", "unban", "ban_all", "unban_all", "delete_from_redis"],
        "airflow_meta": ["clear_dag_runs"],
    }

    if action not in valid_actions.get(entity, []):
        raise ValueError(
            f"The action '{action}' is not valid for entity '{entity}'.\n"
            f"Valid actions for '{entity}' are: {', '.join(valid_actions.get(entity, ['None']))}."
        )
    
    # Validate required parameters for the chosen action
    if entity == "proxy":
        if action in ["ban", "unban"] and not server_identity:
            raise ValueError(f"A 'server_identity' is required for proxy action '{action}'.")
        if action in ["ban", "unban"] and not proxy_url:
            raise ValueError(f"A 'proxy_url' is required for proxy action '{action}'.")
    
    if entity == "account":
        if action in ["ban", "unban"] and not account_id:
            raise ValueError(f"An 'account_id' is required for account action '{action}'.")

    # --- Handle Airflow Meta actions separately as they don't use Thrift ---
    if entity == "airflow_meta":
        dag_id = params.get("dag_id_to_manage")

        if action == "clear_dag_runs":
            clear_scope = params.get("clear_scope")
            logger.info(f"Attempting to delete DagRuns for DAG '{dag_id}' with scope '{clear_scope}'.")

            with create_session() as session:
                dag_run_query = session.query(DagRun).filter(DagRun.dag_id == dag_id)
                
                if clear_scope == "last_run":
                    last_run = dag_run_query.order_by(DagRun.execution_date.desc()).first()
                    if not last_run:
                        logger.info(f"No runs found for DAG '{dag_id}'. Nothing to delete.")
                        print(f"\nNo runs found for DAG '{dag_id}'.\n")
                        return
                    
                    logger.warning(f"Deleting last DagRun for DAG '{dag_id}' (run_id: {last_run.run_id}, execution_date: {last_run.execution_date}). This will also delete its task instances.")
                    # Deleting the DagRun object should cascade and delete related TaskInstances.
                    session.delete(last_run)
                    deleted_count = 1
                else:  # all_runs
                    logger.warning(f"Deleting ALL DagRuns and associated TaskInstances for DAG '{dag_id}'. This will remove all history from the UI.")
                    # To ensure all related data is cleared, we explicitly delete TaskInstances first.
                    # This is safer than relying on DB-level cascades which may not be configured.
                    ti_deleted_count = session.query(TaskInstance).filter(TaskInstance.dag_id == dag_id).delete(synchronize_session=False)
                    logger.info(f"Deleted {ti_deleted_count} TaskInstance records for DAG '{dag_id}'.")
                    
                    deleted_count = dag_run_query.delete(synchronize_session=False)
                
                # The session is committed automatically by the `with create_session()` context manager.
                logger.info(f"Successfully deleted {deleted_count} DagRun(s) for DAG '{dag_id}'.")
                print(f"\nSuccessfully deleted {deleted_count} DagRun(s) for DAG '{dag_id}'.\n")
            return # End execution

    # Handle direct Redis actions separately to avoid creating an unnecessary Thrift connection.
    if action == "delete_from_redis":
        redis_conn_id = params["redis_conn_id"]
        redis_client = _get_redis_client(redis_conn_id)

        if entity == "accounts_and_proxies":
            # --- Delete Proxy ---
            proxy_url = params.get("proxy_url")
            server_identity = params.get("server_identity")

            if proxy_url and server_identity:
                proxy_state_key = f"proxy_status:{server_identity}"

                logger.warning(f"Deleting proxy '{proxy_url}' state from hash '{proxy_state_key}' from Redis.")

                with redis_client.pipeline() as pipe:
                    pipe.hdel(proxy_state_key, proxy_url)
                    results = pipe.execute()
                
                hdel_result = results[0]
                print(f"\nSuccessfully removed proxy '{proxy_url}' from state hash (result: {hdel_result}).")
            else:
                logger.warning("No 'proxy_url' or 'server_identity' provided. Deleting ALL proxy state keys from Redis.")
                patterns = ["proxy_status:*"]
                keys_to_delete = []
                for pattern in patterns:
                    found_keys = [key for key in redis_client.scan_iter(pattern)]
                    if found_keys:
                        logger.info(f"Found {len(found_keys)} keys for pattern '{pattern}'.")
                        keys_to_delete.extend(found_keys)
                    else:
                        logger.info(f"No keys found for pattern '{pattern}'.")

                if not keys_to_delete:
                    print("\nNo proxy keys found to delete.\n")
                else:
                    print(f"\nWARNING: Found {len(keys_to_delete)} proxy-related keys to remove from Redis.")
                    deleted_count = redis_client.delete(*keys_to_delete)
                    print(f"\nSuccessfully removed {deleted_count} proxy-related keys from Redis.\n")

            # --- Delete Account ---
            account_prefix = params.get("account_id")
            pattern = f"account_status:{account_prefix}*" if account_prefix else "account_status:*"
            logger.warning(f"Searching for account status keys in Redis with pattern: '{pattern}'")

            keys_to_delete = [key for key in redis_client.scan_iter(pattern)]

            if not keys_to_delete:
                print(f"\nNo accounts found matching pattern '{pattern}'.\n")
            else:
                print(f"\nWARNING: Found {len(keys_to_delete)} accounts to remove from Redis.")
                for key in keys_to_delete[:10]:
                    print(f"  - {key.decode('utf-8')}")
                if len(keys_to_delete) > 10:
                    print(f"  ... and {len(keys_to_delete) - 10} more.")

                deleted_count = redis_client.delete(*keys_to_delete)
                print(f"\nSuccessfully removed {deleted_count} accounts from Redis.\n")
            
            return # End execution for this action

        if entity == "account":
            account_prefix = params.get("account_id")  # Repurpose account_id param as an optional prefix
            pattern = f"account_status:{account_prefix}*" if account_prefix else "account_status:*"
            logger.warning(f"Searching for account status keys in Redis with pattern: '{pattern}'")

            keys_to_delete = [key for key in redis_client.scan_iter(pattern)]

            if not keys_to_delete:
                print(f"\nNo accounts found matching pattern '{pattern}'.\n")
                return

            print(f"\nWARNING: Found {len(keys_to_delete)} accounts to remove from Redis.")
            for key in keys_to_delete[:10]:
                print(f"  - {key.decode('utf-8')}")
            if len(keys_to_delete) > 10:
                print(f"  ... and {len(keys_to_delete) - 10} more.")

            deleted_count = redis_client.delete(*keys_to_delete)
            print(f"\nSuccessfully removed {deleted_count} accounts from Redis.\n")

        elif entity == "proxy":
            proxy_url = params.get("proxy_url")
            server_identity = params.get("server_identity")

            if proxy_url and server_identity:
                proxy_state_key = f"proxy_status:{server_identity}"

                logger.warning(f"Deleting proxy '{proxy_url}' state from hash '{proxy_state_key}' from Redis.")

                with redis_client.pipeline() as pipe:
                    pipe.hdel(proxy_state_key, proxy_url)
                    results = pipe.execute()
                
                hdel_result = results[0]
                print(f"\nSuccessfully removed proxy '{proxy_url}' from state hash (result: {hdel_result}).\n")
            else:
                logger.warning("No 'proxy_url' or 'server_identity' provided. Deleting ALL proxy state keys from Redis.")
                patterns = ["proxy_status:*"]
                keys_to_delete = []
                for pattern in patterns:
                    found_keys = [key for key in redis_client.scan_iter(pattern)]
                    if found_keys:
                        logger.info(f"Found {len(found_keys)} keys for pattern '{pattern}'.")
                        keys_to_delete.extend(found_keys)
                    else:
                        logger.info(f"No keys found for pattern '{pattern}'.")

                if not keys_to_delete:
                    print("\nNo proxy keys found to delete.\n")
                    return

                print(f"\nWARNING: Found {len(keys_to_delete)} proxy-related keys to remove from Redis.")
                deleted_count = redis_client.delete(*keys_to_delete)
                print(f"\nSuccessfully removed {deleted_count} proxy-related keys from Redis.\n")

        return  # End execution for this action

    client, transport = None, None
    try:
        client, transport = get_thrift_client(host, port)

        if entity == "proxy":
            if action == "list_with_status":
                _list_proxy_statuses(client, server_identity)
            elif action == "ban":
                if not proxy_url: raise ValueError("A 'proxy_url' is required.")
                logger.info(f"Banning proxy '{proxy_url}' for server '{server_identity}'...")
                client.banProxy(proxy_url, server_identity)
                print(f"Successfully sent request to ban proxy '{proxy_url}'.")
            elif action == "unban":
                if not proxy_url: raise ValueError("A 'proxy_url' is required.")
                logger.info(f"Unbanning proxy '{proxy_url}' for server '{server_identity}'...")
                client.unbanProxy(proxy_url, server_identity)
                print(f"Successfully sent request to unban proxy '{proxy_url}'.")
            elif action == "ban_all":
                if server_identity:
                    logger.info(f"Banning all proxies for server '{server_identity}'...")
                    client.banAllProxies(server_identity)
                    print(f"Successfully sent request to ban all proxies for '{server_identity}'.")
                else:
                    logger.info("No server_identity provided. Banning all proxies for ALL servers...")
                    all_statuses = client.getProxyStatus(None)
                    if not all_statuses:
                        print("\nNo proxy statuses found for any server. Nothing to ban.\n")
                        return

                    all_server_identities = sorted(list(set(s.serverIdentity for s in all_statuses)))
                    logger.info(f"Found {len(all_server_identities)} server identities: {all_server_identities}")
                    print(f"Found {len(all_server_identities)} server identities. Sending ban request for each...")

                    success_count = 0
                    fail_count = 0
                    for identity in all_server_identities:
                        try:
                            client.banAllProxies(identity)
                            logger.info(f"  - Sent ban_all for '{identity}'.")
                            success_count += 1
                        except Exception as e:
                            logger.error(f"  - Failed to ban all proxies for '{identity}': {e}")
                            fail_count += 1
                    
                    print(f"\nSuccessfully sent ban_all requests for {success_count} server identities.")
                    if fail_count > 0:
                        print(f"Failed to send ban_all requests for {fail_count} server identities. See logs for details.")
            elif action == "unban_all":
                if server_identity:
                    logger.info(f"Unbanning all proxy statuses for server '{server_identity}'...")
                    client.resetAllProxyStatuses(server_identity)
                    print(f"Successfully sent request to unban all proxy statuses for '{server_identity}'.")
                else:
                    logger.info("No server_identity provided. Unbanning all proxies for ALL servers...")
                    all_statuses = client.getProxyStatus(None)
                    if not all_statuses:
                        print("\nNo proxy statuses found for any server. Nothing to unban.\n")
                        return

                    all_server_identities = sorted(list(set(s.serverIdentity for s in all_statuses)))
                    logger.info(f"Found {len(all_server_identities)} server identities: {all_server_identities}")
                    print(f"Found {len(all_server_identities)} server identities. Sending unban request for each...")

                    success_count = 0
                    fail_count = 0
                    for identity in all_server_identities:
                        try:
                            client.resetAllProxyStatuses(identity)
                            logger.info(f"  - Sent unban_all for '{identity}'.")
                            success_count += 1
                        except Exception as e:
                            logger.error(f"  - Failed to unban all proxies for '{identity}': {e}")
                            fail_count += 1
                    
                    print(f"\nSuccessfully sent unban_all requests for {success_count} server identities.")
                    if fail_count > 0:
                        print(f"Failed to send unban_all requests for {fail_count} server identities. See logs for details.")

        elif entity == "account":
            if action == "list_with_status":
                _list_account_statuses(client, account_id, params["redis_conn_id"])
            elif action == "ban":
                if not account_id: raise ValueError("An 'account_id' is required.")
                reason = f"Manual ban from Airflow mgmt DAG by {socket.gethostname()}"
                logger.info(f"Banning account '{account_id}'...")
                client.banAccount(accountId=account_id, reason=reason)
                print(f"Successfully sent request to ban account '{account_id}'.")
            elif action == "unban":
                if not account_id: raise ValueError("An 'account_id' is required.")
                reason = f"Manual un-ban from Airflow mgmt DAG by {socket.gethostname()}"
                logger.info(f"Unbanning account '{account_id}'...")
                client.unbanAccount(accountId=account_id, reason=reason)
                print(f"Successfully sent request to unban account '{account_id}'.")
            elif action == "unban_all":
                account_prefix = account_id  # Repurpose account_id param as an optional prefix
                logger.info(f"Unbanning all account statuses to ACTIVE (prefix: '{account_prefix or 'ALL'}')...")
                
                all_statuses = client.getAccountStatus(accountId=None, accountPrefix=account_prefix)
                if not all_statuses:
                    print(f"No accounts found with prefix '{account_prefix or 'ALL'}' to unban.")
                    return

                accounts_to_unban = [s.accountId for s in all_statuses]
                logger.info(f"Found {len(accounts_to_unban)} accounts to unban.")
                print(f"Found {len(accounts_to_unban)} accounts. Sending unban request for each...")

                unban_count = 0
                fail_count = 0
                for acc_id in accounts_to_unban:
                    try:
                        reason = f"Manual unban_all from Airflow mgmt DAG by {socket.gethostname()}"
                        client.unbanAccount(accountId=acc_id, reason=reason)
                        logger.info(f"  - Sent unban for '{acc_id}'.")
                        unban_count += 1
                    except Exception as e:
                        logger.error(f"  - Failed to unban account '{acc_id}': {e}")
                        fail_count += 1
                
                print(f"\nSuccessfully sent unban requests for {unban_count} accounts.")
                if fail_count > 0:
                    print(f"Failed to send unban requests for {fail_count} accounts. See logs for details.")
                
                # Optionally, list statuses again to confirm
                print("\n--- Listing statuses after unban_all ---")
                _list_account_statuses(client, account_prefix, params["redis_conn_id"])

        elif entity == "accounts_and_proxies":
            print(f"\n--- Performing action '{action}' on BOTH Proxies and Accounts ---")
            
            # --- Proxy Action ---
            try:
                print("\n-- Running Proxy Action --")
                if action == "list_with_status":
                    _list_proxy_statuses(client, server_identity)
                elif action == "ban":
                    if not proxy_url: raise ValueError("A 'proxy_url' is required.")
                    logger.info(f"Banning proxy '{proxy_url}' for server '{server_identity}'...")
                    client.banProxy(proxy_url, server_identity)
                    print(f"Successfully sent request to ban proxy '{proxy_url}'.")
                elif action == "unban":
                    if not proxy_url: raise ValueError("A 'proxy_url' is required.")
                    logger.info(f"Unbanning proxy '{proxy_url}' for server '{server_identity}'...")
                    client.unbanProxy(proxy_url, server_identity)
                    print(f"Successfully sent request to unban proxy '{proxy_url}'.")
                elif action == "ban_all":
                    if server_identity:
                        logger.info(f"Banning all proxies for server '{server_identity}'...")
                        client.banAllProxies(server_identity)
                        print(f"Successfully sent request to ban all proxies for '{server_identity}'.")
                    else:
                        logger.info("No server_identity provided. Banning all proxies for ALL servers...")
                        all_statuses = client.getProxyStatus(None)
                        if not all_statuses:
                            print("\nNo proxy statuses found for any server. Nothing to ban.\n")
                        else:
                            all_server_identities = sorted(list(set(s.serverIdentity for s in all_statuses)))
                            logger.info(f"Found {len(all_server_identities)} server identities: {all_server_identities}")
                            print(f"Found {len(all_server_identities)} server identities. Sending ban request for each...")

                            success_count = 0
                            fail_count = 0
                            for identity in all_server_identities:
                                try:
                                    client.banAllProxies(identity)
                                    logger.info(f"  - Sent ban_all for '{identity}'.")
                                    success_count += 1
                                except Exception as e:
                                    logger.error(f"  - Failed to ban all proxies for '{identity}': {e}")
                                    fail_count += 1
                            
                            print(f"\nSuccessfully sent ban_all requests for {success_count} server identities.")
                            if fail_count > 0:
                                print(f"Failed to send ban_all requests for {fail_count} server identities. See logs for details.")
                elif action == "unban_all":
                    if server_identity:
                        logger.info(f"Unbanning all proxy statuses for server '{server_identity}'...")
                        client.resetAllProxyStatuses(server_identity)
                        print(f"Successfully sent request to unban all proxy statuses for '{server_identity}'.")
                    else:
                        logger.info("No server_identity provided. Unbanning all proxies for ALL servers...")
                        all_statuses = client.getProxyStatus(None)
                        if not all_statuses:
                            print("\nNo proxy statuses found for any server. Nothing to unban.\n")
                        else:
                            all_server_identities = sorted(list(set(s.serverIdentity for s in all_statuses)))
                            logger.info(f"Found {len(all_server_identities)} server identities: {all_server_identities}")
                            print(f"Found {len(all_server_identities)} server identities. Sending unban request for each...")

                            success_count = 0
                            fail_count = 0
                            for identity in all_server_identities:
                                try:
                                    client.resetAllProxyStatuses(identity)
                                    logger.info(f"  - Sent unban_all for '{identity}'.")
                                    success_count += 1
                                except Exception as e:
                                    logger.error(f"  - Failed to unban all proxies for '{identity}': {e}")
                                    fail_count += 1
                            
                            print(f"\nSuccessfully sent unban_all requests for {success_count} server identities.")
                            if fail_count > 0:
                                print(f"Failed to send unban_all requests for {fail_count} server identities. See logs for details.")
            except Exception as proxy_e:
                logger.error(f"Error during proxy action '{action}': {proxy_e}", exc_info=True)
                print(f"\nERROR during proxy action: {proxy_e}")

            # --- Account Action ---
            try:
                print("\n-- Running Account Action --")
                if action == "list_with_status":
                    _list_account_statuses(client, account_id, params["redis_conn_id"])
                elif action == "ban":
                    if not account_id: raise ValueError("An 'account_id' is required.")
                    reason = f"Manual ban from Airflow mgmt DAG by {socket.gethostname()}"
                    logger.info(f"Banning account '{account_id}'...")
                    client.banAccount(accountId=account_id, reason=reason)
                    print(f"Successfully sent request to ban account '{account_id}'.")
                elif action == "unban":
                    if not account_id: raise ValueError("An 'account_id' is required.")
                    reason = f"Manual un-ban from Airflow mgmt DAG by {socket.gethostname()}"
                    logger.info(f"Unbanning account '{account_id}'...")
                    client.unbanAccount(accountId=account_id, reason=reason)
                    print(f"Successfully sent request to unban account '{account_id}'.")
                elif action == "unban_all":
                    account_prefix = account_id  # Repurpose account_id param as an optional prefix
                    logger.info(f"Unbanning all account statuses to ACTIVE (prefix: '{account_prefix or 'ALL'}')...")
                    
                    all_statuses = client.getAccountStatus(accountId=None, accountPrefix=account_prefix)
                    if not all_statuses:
                        print(f"No accounts found with prefix '{account_prefix or 'ALL'}' to unban.")
                    else:
                        accounts_to_unban = [s.accountId for s in all_statuses]
                        logger.info(f"Found {len(accounts_to_unban)} accounts to unban.")
                        print(f"Found {len(accounts_to_unban)} accounts. Sending unban request for each...")

                        unban_count = 0
                        fail_count = 0
                        for acc_id in accounts_to_unban:
                            try:
                                reason = f"Manual unban_all from Airflow mgmt DAG by {socket.gethostname()}"
                                client.unbanAccount(accountId=acc_id, reason=reason)
                                logger.info(f"  - Sent unban for '{acc_id}'.")
                                unban_count += 1
                            except Exception as e:
                                logger.error(f"  - Failed to unban account '{acc_id}': {e}")
                                fail_count += 1
                        
                        print(f"\nSuccessfully sent unban requests for {unban_count} accounts.")
                        if fail_count > 0:
                            print(f"Failed to send unban requests for {fail_count} accounts. See logs for details.")
                        
                        # Optionally, list statuses again to confirm
                        print("\n--- Listing statuses after unban_all ---")
                        _list_account_statuses(client, account_prefix, params["redis_conn_id"])
            except Exception as account_e:
                logger.error(f"Error during account action '{action}': {account_e}", exc_info=True)
                print(f"\nERROR during account action: {account_e}")

        elif entity == "all":
            if action == "list_with_status":
                print("\nListing all entities...")
                _list_proxy_statuses(client, server_identity)
                _list_account_statuses(client, account_id, params["redis_conn_id"])

    except (PBServiceException, PBUserException) as e:
        logger.error(f"Thrift error performing action '{action}': {e.message}", exc_info=True)
        raise
    except NotImplementedError as e:
        logger.error(f"Feature not implemented: {e}", exc_info=True)
        raise
    except Exception as e:
        logger.error(f"Error performing action '{action}': {e}", exc_info=True)
        raise
    finally:
        if transport and transport.isOpen():
            transport.close()
            logger.info("Thrift connection closed.")

with DAG(
    dag_id="ytdlp_mgmt_proxy_account",
    start_date=days_ago(1),
    schedule=None,
    catchup=False,
    tags=["ytdlp", "mgmt", "master"],
    doc_md="""
    ### YT-DLP Proxy and Account Manager DAG
    This DAG provides tools to manage the state of proxies and accounts used by the `ytdlp-ops-server`.
    Select an `entity` and an `action` to perform.
    """,
    params={
        "management_host": Param(DEFAULT_MANAGEMENT_SERVICE_IP, type="string", title="Management Service Host", description="The hostname or IP of the management service. Can be a Docker container name (e.g., 'envoy-thrift-lb') if on the same network."),
        "management_port": Param(DEFAULT_MANAGEMENT_SERVICE_PORT, type="integer", title="Management Service Port", description="The port of the dedicated management service."),
        "entity": Param(
            "accounts_and_proxies",
            type="string",
            enum=["account", "proxy", "accounts_and_proxies", "airflow_meta"],
            description="The type of entity to manage.",
        ),
        "action": Param(
            "list_with_status",
            type="string",
            enum=["list_with_status", "ban", "unban", "ban_all", "unban_all", "delete_from_redis", "clear_dag_runs"],
            description="""The management action to perform.
            ---
            #### Actions for `entity: proxy`
            - `list_with_status`: View status of all proxies, optionally filtered by `server_identity`.
            - `ban`: Ban a specific proxy for a given `server_identity`. Requires `proxy_url`.
            - `unban`: Un-ban a specific proxy. Requires `proxy_url`.
            - `ban_all`: Sets the status of all proxies for a given `server_identity` (or all servers) to `BANNED`.
            - `unban_all`: Resets the status of all proxies for a given `server_identity` (or all servers) to `ACTIVE`.
            - `delete_from_redis`: **(Destructive)** Deletes proxy **state** from Redis. This action does not remove the proxy from the service's configuration, but rather resets its status (ban/active, success/failure counts) to the default. The service will continue to manage the proxy. If `proxy_url` and `server_identity` are provided, it deletes a single proxy's state. If they are omitted, it deletes **ALL** proxy state keys (`proxy_status:*`).

            #### Actions for `entity: account`
            - `list_with_status`: View status of all accounts, optionally filtered by `account_id` (as a prefix).
            - `ban`: Ban a specific account. Requires `account_id`.
            - `unban`: Un-ban a specific account. Requires `account_id`.
            - `unban_all`: Sets the status of all accounts (or those matching a prefix in `account_id`) to `ACTIVE`.
            - `delete_from_redis`: **(Destructive)** Deletes account status keys from Redis. This permanently removes the account from being tracked by the system. This is different from `unban`. Use with caution.

            #### Actions for `entity: accounts_and_proxies`
            - This entity performs the selected action on **both** proxies and accounts where applicable.
            - `list_with_status`: View statuses for both proxies and accounts.
            - `ban`: Ban a specific proxy AND a specific account. Requires `proxy_url`, `server_identity`, and `account_id`.
            - `unban`: Un-ban a specific proxy AND a specific account. Requires `proxy_url`, `server_identity`, and `account_id`.
            - `ban_all`: Ban all proxies for a `server_identity` (or all servers). Does not affect accounts.
            - `unban_all`: Un-ban all proxies for a `server_identity` (or all servers) AND all accounts (optionally filtered by `account_id` as a prefix).
            - `delete_from_redis`: Deletes proxy and account **state** from Redis. For proxies, this resets their status but they remain managed by the service. For accounts, this permanently removes them from the system's tracking. If `proxy_url` and `server_identity` are provided, it deletes a single proxy's state. If they are omitted, it deletes **ALL** proxy state (keys matching `proxy_status:*`). It will also delete all accounts matching the `account_id` prefix (or all accounts if `account_id` is empty).

            #### Actions for `entity: airflow_meta`
            - `clear_dag_runs`: **(Destructive)** Deletes DAG run history and associated task instances from the database, removing them from the UI. This allows the runs to be re-created if backfilling is enabled.
                - `clear_scope: last_run`: Deletes only the most recent DAG run and its task instances.
                - `clear_scope: all_runs`: Deletes all historical DAG runs and task instances for the selected DAG.
            """,
        ),
        "server_identity": Param(
            None,
            type=["null", "string"],
            description="The identity of the server instance (for proxy management). Leave blank to list all.",
        ),
        "proxy_url": Param(
            None,
            type=["null", "string"],
            description="The proxy URL to act upon (e.g., 'socks5://host:port').",
        ),
        "account_id": Param(
            None,
            type=["null", "string"],
            description="The account ID to act upon. For `unban_all` or `delete_from_redis` on accounts, this can be an optional prefix.",
        ),
        "redis_conn_id": Param(
            DEFAULT_REDIS_CONN_ID,
            type="string",
            title="Redis Connection ID",
            description="The Airflow connection ID for the Redis server (used for 'delete_from_redis' and for fetching detailed account status).",
        ),
        "dag_id_to_manage": Param(
            "ytdlp_ops_worker_per_url",
            type="string",
            enum=["ytdlp_ops_worker_per_url", "ytdlp_ops_orchestrator"],
            title="[Airflow Meta] DAG ID",
            description="The DAG ID to perform the action on.",
        ),
        "clear_scope": Param(
            "last_run",
            type="string",
            enum=["last_run", "all_runs"],
            title="[Airflow Meta] Clear Scope",
            description="For 'clear_dag_runs' action, specifies the scope of runs to clear.",
        ),
    },
) as dag:
    system_management_task = PythonOperator(
        task_id="system_management_task",
        python_callable=manage_system_callable,
    )
