#!/usr/bin/env python3
"""
CLI tool to orchestrate multi-stage profile simulations.
"""
import argparse
import logging
import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from types import SimpleNamespace

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

# Import the main functions from the tools we are wrapping
from .profile_setup_tool import main_setup_profiles
from .stress_policy_tool import main_stress_policy

logger = logging.getLogger(__name__)

# Define default policy paths relative to the project root
PROJECT_ROOT = Path(__file__).resolve().parent.parent
POLICY_DIR = PROJECT_ROOT / 'policies'
POLICY_FILE_SETUP = str(POLICY_DIR / '6_simulation_policy.yaml')
POLICY_FILE_AUTH = str(POLICY_DIR / '7_continuous_auth.yaml')
POLICY_FILE_DOWNLOAD = str(POLICY_DIR / '8_continuous_download.yaml')

def add_simulation_parser(subparsers):
    """Adds the parser for the 'simulation' command."""
    parser = subparsers.add_parser(
        'simulation',
        description="Run multi-stage profile simulations (setup, auth, download). This provides a unified interface for the simulation workflow.",
        formatter_class=argparse.RawTextHelpFormatter,
        help="Run multi-stage profile simulations."
    )
    
    # Common arguments for all simulation subcommands
    common_parser = argparse.ArgumentParser(add_help=False)
    common_parser.add_argument('--env-file', default=None, help="Path to a .env file to load. Overrides setting from policy file.")
    common_parser.add_argument('--redis-host', default=None, help='Redis host. Overrides policy and .env file.')
    common_parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Overrides policy and .env file.')
    common_parser.add_argument('--redis-password', default=None, help='Redis password. Overrides policy and .env file.')
    common_parser.add_argument('--env', default='sim', help="Environment name for Redis key prefix. Default: 'sim'.")
    common_parser.add_argument('--expire-time-shift-minutes', type=int, default=None, help="Consider URLs expiring in N minutes as expired. Overrides policy.")
    common_parser.add_argument('--verbose', action='store_true', help="Enable verbose logging.")
    
    sim_subparsers = parser.add_subparsers(dest='simulation_command', help='Simulation stage to run', required=True)

    # --- Setup ---
    setup_parser = sim_subparsers.add_parser('setup', help='Set up profiles for a simulation.', parents=[common_parser])
    setup_parser.add_argument('--policy-file', dest='policy', default=POLICY_FILE_SETUP, help=f'Path to the setup policy YAML file. Default: {POLICY_FILE_SETUP}')
    setup_parser.add_argument('--preserve-profiles', action='store_true', help="Do not clean up existing profiles.")
    setup_parser.add_argument('--reset-global-counters', action='store_true', help="Reset global counters like 'failed_lock_attempts'.")

    # --- Auth ---
    auth_parser = sim_subparsers.add_parser('auth', help='Run the authentication (get-info) part of the simulation.', parents=[common_parser])
    auth_parser.add_argument('--policy-file', dest='policy', default=POLICY_FILE_AUTH, help=f'Path to the auth simulation policy file. Default: {POLICY_FILE_AUTH}')
    auth_parser.add_argument('--set', action='append', default=[], help="Override a policy setting using 'key.subkey=value' format.")

    # --- Download ---
    download_parser = sim_subparsers.add_parser('download', help='Run the download part of the simulation.', parents=[common_parser])
    download_parser.add_argument('--policy-file', dest='policy', default=POLICY_FILE_DOWNLOAD, help=f'Path to the download simulation policy file. Default: {POLICY_FILE_DOWNLOAD}')
    download_parser.add_argument('--set', action='append', default=[], help="Override a policy setting using 'key.subkey=value' format.")

def main_simulation(args):
    """Main dispatcher for 'simulation' command."""
    # --- Load policy to get simulation parameters ---
    policy = {}
    # The 'policy' attribute is guaranteed to exist by the arg parser for all subcommands
    if not yaml:
        logger.error("Cannot load policy file because PyYAML is not installed.")
        return 1
    try:
        with open(args.policy, 'r') as f:
            # We only need the first document if it's a multi-policy file
            policy = yaml.safe_load(f) or {}
    except (IOError, yaml.YAMLError) as e:
        logger.error(f"Failed to load or parse policy file {args.policy}: {e}")
        return 1

    sim_params = policy.get('simulation_parameters', {})
    effective_env_file = args.env_file or sim_params.get('env_file')

    if args.simulation_command == 'setup':
        # Create an args object that main_setup_profiles expects
        setup_args = SimpleNamespace(
            policy_file=args.policy,
            env_file=effective_env_file,
            preserve_profiles=args.preserve_profiles,
            reset_global_counters=args.reset_global_counters,
            verbose=args.verbose,
            redis_host=args.redis_host,
            redis_port=args.redis_port,
            redis_password=args.redis_password
        )
        return main_setup_profiles(setup_args)

    elif args.simulation_command == 'auth':
        # This command runs the stress tool in auth (fetch_only) mode.
        # It is expected that the policy-enforcer is run as a separate process.
        stress_args = SimpleNamespace(
            policy=args.policy, policy_name=None, list_policies=False, show_overrides=False,
            set=args.set, profile_prefix=None, start_from_url_index=None, auto_merge_fragments=None,
            remove_fragments_after_merge=None, fragments_dir=None, remote_dir=None, cleanup=None,
            verbose=args.verbose, dry_run=False, disable_log_writing=False,
            # Redis connection args
            env_file=effective_env_file, redis_host=args.redis_host, redis_port=args.redis_port,
            redis_password=args.redis_password, env=args.env, key_prefix=None,
            expire_time_shift_minutes=args.expire_time_shift_minutes
        )

        logger.info("\n--- Starting Auth Simulation (stress-policy) ---")
        return main_stress_policy(stress_args)
    
    elif args.simulation_command == 'download':
        # This is simpler, just runs the stress tool in download mode.
        stress_args = SimpleNamespace(
            policy=args.policy, policy_name=None, list_policies=False, show_overrides=False,
            set=args.set, profile_prefix=None, start_from_url_index=None, auto_merge_fragments=None,
            remove_fragments_after_merge=None, fragments_dir=None, remote_dir=None, cleanup=None,
            verbose=args.verbose, dry_run=False, disable_log_writing=False,
            # Redis connection args
            env_file=effective_env_file, redis_host=args.redis_host, redis_port=args.redis_port,
            redis_password=args.redis_password, env=args.env, key_prefix=None,
            expire_time_shift_minutes=args.expire_time_shift_minutes
        )
        logger.info("\n--- Starting Download Simulation (stress-policy) ---")
        return main_stress_policy(stress_args)
    
    return 1 # Should not be reached