#!/usr/bin/env python3
"""
Client script to manage proxies in the YTTokenOpService.

This script allows you to list, ban, unban, and reset proxies that are managed
by a ytdlp-ops-server instance via Redis.
"""

import argparse
import sys
import os
import logging
from pathlib import Path
import datetime

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



try:
    from thrift.transport import TTransport
    from pangramia.yt.exceptions.ttypes import PBServiceException, PBUserException
    from pangramia.yt.common.constants import ErrorCode
    from tabulate import tabulate
    from yt_ops_services.client_utils import get_thrift_client, format_timestamp
except ImportError as e:
    print(f"Error importing required modules: {e}")
    print("Please ensure you have installed dependencies by running: pip install -e .")
    sys.exit(1)




def main():
    parser = argparse.ArgumentParser(
        description="Manage proxies for the YTDLP Operations Server.\n\n"
                    "This script allows you to list, ban, unban, and reset proxies that are managed\n"
                    "by a ytdlp-ops-server instance via Redis. It provides a command-line interface\n"
                    "to interact with the proxy management features of the server.",
        epilog="Usage examples:\n"
               "  # List statuses for a specific server identity\n"
               "  python proxy_manager_client.py list --server-identity ytdlp-ops-airflow-service\n\n"
               "  # Ban a proxy for a specific server\n"
               "  python proxy_manager_client.py ban --server-identity ytdlp-ops-airflow-service --proxy-url socks5://proxy.example.com:1080\n\n"
               "  # Unban a proxy\n"
               "  python proxy_manager_client.py unban --server-identity ytdlp-ops-airflow-service --proxy-url socks5://proxy.example.com:1080\n\n"
               "  # Reset all proxies for a server to ACTIVE\n"
               "  python proxy_manager_client.py reset --server-identity ytdlp-ops-airflow-service",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument('--host', default=os.getenv('YTDLP_HOST', '127.0.0.1'), help="Server host (default: 127.0.0.1 or YTDLP_HOST env). Using 127.0.0.1 avoids harmless connection errors when the local Envoy proxy only listens on IPv4.")
    parser.add_argument('--port', type=int, default=int(os.getenv('YTDLP_PORT', '9080')), help='Server port (default: 9080 or YTDLP_PORT env)')

    subparsers = parser.add_subparsers(dest='command', required=True, help='Available commands')

    # List command
    list_parser = subparsers.add_parser(
        'list',
        help='List proxy statuses for a given server identity.',
        description="List the status of all proxies associated with a specific server identity.\n"
                    "The status includes:\n"
                    "- Server: The server identity.\n"
                    "- Proxy URL: The URL of the proxy.\n"
                    "- Status: ACTIVE or BANNED.\n"
                    "- Success: Count of successful uses.\n"
                    "- Failures: Count of failed uses.\n"
                    "- Last Success: Timestamp of the last successful use.\n"
                    "- Last Failure: Timestamp of the last failed use.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    list_parser.add_argument('--server-identity', type=str, help='The identity of the server to query. If not provided, shows status for the connected server instance.')

    # Ban command
    ban_parser = subparsers.add_parser(
        'ban',
        help='Ban a specific proxy for a server.',
        description="Manually set a proxy's status to BANNED for a specific server identity.\n"
                    "A banned proxy will not be used for future requests by that server instance\n"
                    "until it is unbanned or reset.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    ban_parser.add_argument('--server-identity', type=str, required=True, help='The identity of the server where the proxy should be banned.')
    ban_parser.add_argument('--proxy-url', type=str, required=True, help="The full URL of the proxy to ban (e.g., 'socks5://host:port').")

    # Unban command
    unban_parser = subparsers.add_parser(
        'unban',
        help='Unban a specific proxy for a server.',
        description="Manually set a proxy's status to ACTIVE for a specific server identity.\n"
                    "This will allow the server instance to use the proxy for future requests.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    unban_parser.add_argument('--server-identity', type=str, required=True, help='The identity of the server where the proxy should be unbanned.')
    unban_parser.add_argument('--proxy-url', type=str, required=True, help="The full URL of the proxy to unban (e.g., 'socks5://host:port').")

    # Reset command
    reset_parser = subparsers.add_parser(
        'reset',
        help='Reset all proxy statuses for a server to ACTIVE.',
        description="Reset the status of all proxies associated with a specific server identity to ACTIVE.\n"
                    "This is useful for clearing all bans and making all configured proxies available again.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    reset_parser.add_argument('--server-identity', type=str, required=True, help='The identity of the server whose proxies should be reset.')

    args = parser.parse_args()

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

        if args.command == 'list':
            logger.info(f"Getting proxy statuses for server: {args.server_identity or 'local server'}")
            statuses = client.getProxyStatus(args.server_identity)
            if not statuses:
                print("\nThe server reported no proxy statuses.")
                print("This can happen if no proxies are configured, or if all configured proxies failed their initial health check on server startup.\n")
            else:
                # Determine which proxy is next in rotation for each server identity
                next_proxies = {s.serverIdentity: s.proxyUrl for s in statuses if '(next)' in s.status}

                status_list = []
                for s in statuses:
                    is_next = next_proxies.get(s.serverIdentity) == s.proxyUrl
                    status_list.append({
                        "Server": s.serverIdentity,
                        "Proxy URL": f"{s.proxyUrl} ->" if is_next else s.proxyUrl,
                        "Status": s.status.replace(" (next)", ""),
                        "Success": s.successCount,
                        "Failures": s.failureCount,
                        "Last Success": format_timestamp(s.lastSuccessTimestamp),
                        "Last Failure": format_timestamp(s.lastFailureTimestamp),
                    })
                print("\n--- Proxy Statuses ---")
                print(tabulate(status_list, headers="keys", tablefmt="grid"))
                print("----------------------\n")

        elif args.command == 'ban':
            logger.info(f"Banning proxy '{args.proxy_url}' for server '{args.server_identity}'...")
            success = client.banProxy(args.proxy_url, args.server_identity)
            if success:
                print(f"Successfully banned proxy '{args.proxy_url}' for server '{args.server_identity}'.")
            else:
                print("Failed to ban proxy. Check server logs for details.")
                sys.exit(1)

        elif args.command == 'unban':
            logger.info(f"Unbanning proxy '{args.proxy_url}' for server '{args.server_identity}'...")
            success = client.unbanProxy(args.proxy_url, args.server_identity)
            if success:
                print(f"Successfully unbanned proxy '{args.proxy_url}' for server '{args.server_identity}'.")
            else:
                print("Failed to unban proxy. Check server logs for details.")
                sys.exit(1)

        elif args.command == 'reset':
            logger.info(f"Resetting all proxy statuses for server '{args.server_identity}'...")
            success = client.resetAllProxyStatuses(args.server_identity)
            if success:
                print(f"Successfully reset all proxy statuses for server '{args.server_identity}'.")
            else:
                print("Failed to reset all proxy statuses. Check server logs for details.")
                sys.exit(1)

    except (PBServiceException, PBUserException) as e:
        if hasattr(e, 'errorCode') and e.errorCode == ErrorCode.NOT_IMPLEMENTED:
            logger.error(f"Action '{args.command}' is not implemented by the server. It may be running in the wrong service mode.")
            print(f"Error: The server does not support the action '{args.command}'.")
            print("Please check that the server is running in 'all-in-one' or 'management' mode.")
        else:
            logger.error(f"Thrift error performing action '{args.command}': {e.message}", exc_info=True)
            print(f"Error: {e.message}")
        sys.exit(1)
    except TTransport.TTransportException as e:
        # The logger.error is not needed here because TSocket already logs connection errors.
        print(f"Error: Connection to server at {args.host}:{args.port} failed. Is the server running?")
        print(f"Details: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"An unexpected error occurred: {e}", exc_info=True)
        print(f"An unexpected error occurred: {e}")
        sys.exit(1)
    finally:
        if transport and transport.isOpen():
            transport.close()
            logger.info("Thrift connection closed.")


if __name__ == "__main__":
    main()
