#!/usr/bin/env python3
"""
CLI tool to set up profiles from a YAML policy file.
"""

import argparse
import json
import logging
import os
import subprocess
import sys
from typing import List

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

logger = logging.getLogger(__name__)

def run_command(cmd: List[str], capture: bool = False):
    """Runs a command and raises an exception on failure."""
    logger.debug(f"Running command: {' '.join(cmd)}")
    # check=True will raise CalledProcessError on non-zero exit codes
    result = subprocess.run(
        cmd, 
        capture_output=capture, 
        text=True,
        check=True
    )
    return result

def add_setup_profiles_parser(subparsers):
    """Adds the parser for the 'setup-profiles' command."""
    parser = subparsers.add_parser(
        'setup-profiles',
        description="Set up profiles for a simulation or test run based on a policy file.",
        formatter_class=argparse.RawTextHelpFormatter,
        help="Set up profiles from a policy file."
    )
    parser.add_argument('--policy', '--policy-file', dest='policy_file', required=True, help="Path to the YAML profile setup policy file.")
    parser.add_argument('--env', help="Override the environment name from the policy file. For multi-setup files, this will override the 'env' for ALL setups being run.")
    parser.add_argument('--env-file', help="Override the env_file setting in the policy.")
    parser.add_argument('--auth-only', action='store_true', help='In a multi-setup policy file, run only the auth_profile_setup.')
    parser.add_argument('--download-only', action='store_true', help='In a multi-setup policy file, run only the download_profile_setup.')
    parser.add_argument('--redis-host', default=None, help='Redis host. Overrides policy and .env file.')
    parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Overrides policy and .env file.')
    parser.add_argument('--redis-password', default=None, help='Redis password. Overrides policy and .env file.')
    parser.add_argument('--preserve-profiles', action='store_true', help="Do not clean up existing profiles; create only what is missing.")
    parser.add_argument('--cleanup-prefix', action='append', help="Prefix of profiles to delete before setup. Can be specified multiple times. Overrides policy-based cleanup.")
    parser.add_argument('--cleanup-all', action='store_true', help="(Destructive) Delete ALL data for the environment (profiles, proxies, counters) before setup. Overrides all other cleanup options.")
    parser.add_argument('--reset-global-counters', action='store_true', help="Reset global counters like 'failed_lock_attempts'.")
    parser.add_argument('--verbose', action='store_true', help="Enable verbose logging.")
    return parser

def _run_setup_for_env(profile_setup: dict, common_args: list, args: argparse.Namespace) -> int:
    """Runs the profile setup logic for a given configuration block."""
    if args.cleanup_all:
        logger.info("--- (DESTRUCTIVE) Cleaning up all data for the environment via --cleanup-all ---")
        try:
            cleanup_cmd = ['bin/ytops-client', 'profile', 'delete-all', '--confirm'] + common_args
            run_command(cleanup_cmd)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to clean up all data for the environment: {e}")
            return 1 # Stop if cleanup fails
        
        # Disable other cleanup logic
        profile_setup['cleanup_before_run'] = False
        args.cleanup_prefix = None
    
    # If --cleanup-prefix is provided, it takes precedence over policy settings
    if args.cleanup_prefix:
        logger.info("--- Cleaning up profiles based on --cleanup-prefix ---")
        for prefix in args.cleanup_prefix:
            try:
                list_cmd = ['bin/ytops-client', 'profile', 'list', '--format', 'json'] + common_args
                result = run_command(list_cmd, capture=True)
                profiles_to_delete = [p for p in json.loads(result.stdout) if p['name'].startswith(prefix)]
                
                if not profiles_to_delete:
                    logger.info(f"No profiles with prefix '{prefix}' found to delete.")
                    continue

                logger.info(f"Found {len(profiles_to_delete)} profiles with prefix '{prefix}' to delete.")
                for p in profiles_to_delete:
                    delete_cmd = ['bin/ytops-client', 'profile', 'delete', p['name'], '--confirm'] + common_args
                    run_command(delete_cmd)
            except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
                logger.warning(f"Could not list or parse existing profiles with prefix '{prefix}' for cleanup. Error: {e}")
        
        # Disable policy-based cleanup as we've handled it via CLI
        profile_setup['cleanup_before_run'] = False

    if args.preserve_profiles:
        if profile_setup.get('cleanup_before_run'):
            logger.info("--preserve-profiles is set, overriding 'cleanup_before_run: true' from policy.")
            profile_setup['cleanup_before_run'] = False

    if args.reset_global_counters:
        logger.info("--- Resetting global counters ---")
        try:
            reset_cmd = ['bin/ytops-client', 'profile', 'reset-global-counters'] + common_args
            run_command(reset_cmd)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to reset global counters: {e}")

    if profile_setup.get('cleanup_before_run'):
        logger.info("--- Cleaning up old profiles ---")
        for pool in profile_setup.get('pools', []):
            prefix = pool.get('prefix')
            if prefix:
                try:
                    list_cmd = ['bin/ytops-client', 'profile', 'list', '--format', 'json'] + common_args
                    result = run_command(list_cmd, capture=True)
                    profiles = [p for p in json.loads(result.stdout) if p['name'].startswith(prefix)]
                    
                    if not profiles:
                        logger.info(f"No profiles with prefix '{prefix}' found to delete.")
                        continue

                    logger.info(f"Found {len(profiles)} profiles with prefix '{prefix}' to delete.")
                    for p in profiles:
                        delete_cmd = ['bin/ytops-client', 'profile', 'delete', p['name'], '--confirm'] + common_args
                        run_command(delete_cmd)
                except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
                    logger.warning(f"Could not list existing profiles with prefix '{prefix}' for cleanup. Assuming none exist. Error: {e}")

    existing_profiles = set()
    if not profile_setup.get('cleanup_before_run'):
        logger.info("--- Checking for existing profiles ---")
        try:
            list_cmd = ['bin/ytops-client', 'profile', 'list', '--format', 'json'] + common_args
            result = run_command(list_cmd, capture=True)
            profiles_data = json.loads(result.stdout)
            existing_profiles = {p['name'] for p in profiles_data}
            logger.info(f"Found {len(existing_profiles)} existing profiles.")
        except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
            logger.error(f"Failed to list existing profiles. Will attempt to create all profiles. Error: {e}")

    logger.info("--- Creating new profiles (if needed) ---")
    for pool in profile_setup.get('pools', []):
        prefix = pool.get('prefix')
        proxy = pool.get('proxy')
        count = pool.get('count', 0)
        start_in_rest_minutes = pool.get('start_in_rest_for_minutes')

        profiles_in_pool_created = 0
        # Get a list of all profiles that should exist for this pool
        profile_names_in_pool = [f"{prefix}_{i}" for i in range(count)]

        for profile_name in profile_names_in_pool:
            if profile_name in existing_profiles:
                logger.debug(f"Profile '{profile_name}' already exists, preserving.")
                continue

            try:
                create_cmd = ['bin/ytops-client', 'profile', 'create', profile_name, proxy] + common_args
                run_command(create_cmd)
                profiles_in_pool_created += 1
            except subprocess.CalledProcessError as e:
                logger.error(f"Failed to create profile '{profile_name}': {e}")
        
        if profiles_in_pool_created > 0:
            logger.info(f"Created {profiles_in_pool_created} new profile(s) for pool '{prefix}'.")
        elif count > 0:
            logger.info(f"No new profiles needed for pool '{prefix}'. All {count} profile(s) already exist.")

        # If requested, put the proxy for this pool into a RESTING state.
        # This is done even for existing profiles when --preserve-profiles is used.
        if start_in_rest_minutes and proxy:
            logger.info(f"Setting proxy '{proxy}' for pool '{prefix}' to start in RESTING state for {start_in_rest_minutes} minutes.")
            try:
                # The 'set-proxy-state' command takes '--duration-minutes'
                set_state_cmd = ['bin/ytops-client', 'profile', 'set-proxy-state', proxy, 'RESTING',
                                 '--duration-minutes', str(start_in_rest_minutes)] + common_args
                run_command(set_state_cmd)
            except subprocess.CalledProcessError as e:
                logger.error(f"Failed to set initial REST state for proxy '{proxy}': {e}")
    
    return 0


def main_setup_profiles(args):
    """Main logic for the 'setup-profiles' command."""
    if not yaml: return 1

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        logger.setLevel(logging.DEBUG)
    else:
        if not logging.getLogger().handlers:
            logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        logger.setLevel(logging.INFO)

    try:
        with open(args.policy_file, 'r', encoding='utf-8') 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

    # --- Expand common_pools if defined ---
    common_pools_def = policy.get('common_pools')
    if common_pools_def:
        logger.info("Found 'common_pools' definition. Expanding into setup blocks.")
        expanded_pools = []
        for pool_def in common_pools_def:
            # Handle both 'prefix' (string) and 'prefixes' (list) for flexibility
            prefixes = pool_def.get('prefixes', [])
            if 'prefix' in pool_def and pool_def['prefix'] not in prefixes:
                prefixes.append(pool_def['prefix'])
            
            if not prefixes:
                logger.warning(f"A pool in 'common_pools' is missing a 'prefix' or 'prefixes' key. Skipping: {pool_def}")
                continue

            for p in prefixes:
                new_pool = pool_def.copy()
                new_pool['prefix'] = p
                if 'prefixes' in new_pool: del new_pool['prefixes']
                expanded_pools.append(new_pool)
        
        for setup_key in ['auth_profile_setup', 'download_profile_setup']:
            if setup_key in policy and policy[setup_key].get('use_common_pools'):
                logger.debug(f"Applying {len(expanded_pools)} common pool definitions to '{setup_key}'.")
                policy[setup_key]['pools'] = expanded_pools

    sim_params = policy.get('simulation_parameters', {})
    setups_to_run = []

    if not args.download_only and 'auth_profile_setup' in policy:
        setups_to_run.append(('Auth', policy['auth_profile_setup']))
    
    if not args.auth_only and 'download_profile_setup' in policy:
        setups_to_run.append(('Download', policy['download_profile_setup']))

    # Backward compatibility for old single-block format
    if not setups_to_run and 'profile_setup' in policy:
        legacy_config = policy['profile_setup']
        # Synthesize the env from the global section for the legacy block
        legacy_config['env'] = args.env or sim_params.get('env')
        setups_to_run.append(('Legacy', legacy_config))

    if not setups_to_run:
        logger.error("No 'auth_profile_setup', 'download_profile_setup', or legacy 'profile_setup' block found in policy file.")
        return 1

    env_file = args.env_file or sim_params.get('env_file')
    if load_dotenv:
        if load_dotenv(env_file):
            print(f"Loaded environment variables from {env_file or '.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

    base_common_args = []
    if env_file: base_common_args.extend(['--env-file', env_file])
    
    redis_host = args.redis_host or os.getenv('REDIS_HOST') or os.getenv('MASTER_HOST_IP') or sim_params.get('redis_host')
    if redis_host: base_common_args.extend(['--redis-host', redis_host])
    
    redis_port = args.redis_port
    if redis_port is None:
        redis_port_env = os.getenv('REDIS_PORT')
        redis_port = int(redis_port_env) if redis_port_env and redis_port_env.isdigit() else sim_params.get('redis_port')
    if redis_port: base_common_args.extend(['--redis-port', str(redis_port)])
    
    redis_password = args.redis_password or os.getenv('REDIS_PASSWORD') or sim_params.get('redis_password')
    if redis_password: base_common_args.extend(['--redis-password', redis_password])

    if args.verbose: base_common_args.append('--verbose')

    for setup_name, setup_config in setups_to_run:
        logger.info(f"--- Running setup for {setup_name} simulation ---")
        
        effective_env = args.env or setup_config.get('env')
        if not effective_env:
            logger.error(f"Could not determine environment for '{setup_name}' setup. Please specify 'env' in the policy block or via --env.")
            return 1
        
        env_common_args = base_common_args + ['--env', effective_env]
        
        if _run_setup_for_env(setup_config, env_common_args, args) != 0:
            return 1

    logger.info("\n--- All profile setups complete. ---")
    logger.info("You can now run the policy enforcer to manage the profiles:")
    logger.info("e.g., bin/ytops-client policy-enforcer --policy-file policies/8_unified_simulation_enforcer.yaml --live")
    return 0