#!/usr/bin/env python3
"""
AWS Step Functions Execution Retry Tool

This program takes a list of execution names and starts new executions for each one
with an incremented retry count. The new execution names follow the pattern:
<original-execution-name>-retry-<retry-count>

The retry count is automatically determined by examining existing executions with
similar names and incrementing from the highest found retry count.

If the input execution name is already a retry (e.g., "ProcessData-001-retry-2"),
the script will extract the base name ("ProcessData-001") and create the next
retry in the sequence (e.g., "ProcessData-001-retry-3").

Usage:
    python bulk_start_new_execution.py <state_machine_name> --executions <execution1> <execution2> ... [options]
    python bulk_start_new_execution.py <state_machine_name> --file <executions_file> [options]
    
Examples:
    python bulk_start_new_execution.py MyStateMachine --executions ProcessData-001 ProcessData-002
    python bulk_start_new_execution.py MyStateMachine --executions ProcessData-001-retry-1 --dry-run
    python bulk_start_new_execution.py MyStateMachine --file executions.txt --dry-run
    python bulk_start_new_execution.py MyStateMachine --executions ProcessData-001 --dry-run
"""

import boto3
import json
import sys
import os
import argparse
import re
from datetime import datetime, timezone
from botocore.exceptions import ClientError, NoCredentialsError
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()


def setup_aws_session():
    """
    Setup AWS session with proper error handling for credentials.
    Returns boto3 Step Functions client.
    """
    try:
        # Check if .env file exists and report what we found
        env_file_path = os.path.join(os.getcwd(), '.env')
        if os.path.exists(env_file_path):
            print("✓ Found .env file - loading AWS credentials")
        
        # Create session - this will use environment variables or AWS CLI config
        session = boto3.Session()
        
        # Test credentials by getting caller identity
        sts = session.client('sts')
        identity = sts.get_caller_identity()
        print(f"✓ Connected to AWS as: {identity.get('Arn', 'Unknown')}")
        print(f"✓ Account ID: {identity.get('Account', 'Unknown')}")
        
        # Show current region
        region = session.region_name or os.environ.get('AWS_DEFAULT_REGION', 'Unknown')
        print(f"✓ Region: {region}")
        
        # Create Step Functions client
        sfn_client = session.client('stepfunctions')
        return sfn_client
        
    except NoCredentialsError:
        print("❌ Error: No AWS credentials found!")
        sys.exit(1)
        
    except ClientError as e:
        print(f"❌ Error connecting to AWS: {e}")
        sys.exit(1)


def find_state_machine_arn(sfn_client, state_machine_name):
    """
    Find the ARN of a state machine by name.
    """
    try:
        # List all state machines and find the one with matching name
        response = sfn_client.list_state_machines()
        
        for sm in response.get('stateMachines', []):
            if sm['name'] == state_machine_name:
                return sm['stateMachineArn']
        
        # If not found, print available state machines
        print(f"❌ State machine '{state_machine_name}' not found!")
        print("\nAvailable state machines:")
        for sm in response.get('stateMachines', []):
            print(f"  - {sm['name']}")
        
        return None
        
    except ClientError as e:
        print(f"❌ Error listing state machines: {e}")
        return None


def find_highest_retry_count(sfn_client, state_machine_arn, execution_name):
    """
    Find the highest retry count for executions with the given execution name.
    If the execution name is already a retry, extract the base name first.
    Returns the base execution name and the next retry count to use.
    """
    try:
        # Check if the execution name is already a retry
        retry_pattern = re.compile(r"^(.+)-retry-(\d+)$")
        match = retry_pattern.match(execution_name)
        
        if match:
            # This is already a retry execution, extract the base name
            base_execution_name = match.group(1)
            current_retry_count = int(match.group(2))
            print(f"   Detected retry execution: base='{base_execution_name}', current_retry={current_retry_count}")
        else:
            # This is an original execution name
            base_execution_name = execution_name
            current_retry_count = -1  # Will be incremented to 0
        
        # List recent executions to find existing retry counts for the base name
        response = sfn_client.list_executions(
            stateMachineArn=state_machine_arn,
            maxResults=1000  # Get a good sample to find retry patterns
        )
        
        executions = response.get('executions', [])
        highest_retry = current_retry_count  # Start from current retry count if it's a retry
        
        # Pattern to match retry executions for the base name: base-name-retry-N
        base_retry_pattern = re.compile(rf"^{re.escape(base_execution_name)}-retry-(\d+)$")
        
        for execution in executions:
            exec_name = execution.get('name', '')
            match = base_retry_pattern.match(exec_name)
            if match:
                retry_count = int(match.group(1))
                highest_retry = max(highest_retry, retry_count)
        
        # Return the base name and next retry count (starts at 0 if no retries found)
        next_retry = highest_retry + 1
        return base_execution_name, next_retry
        
    except ClientError as e:
        print(f"⚠️  Warning: Could not check existing retry counts for {execution_name}: {e}")
        # Default to treating as base name with retry-0 if we can't check
        base_name = execution_name.split('-retry-')[0] if '-retry-' in execution_name else execution_name
        return base_name, 0


def get_execution_input(sfn_client, state_machine_arn, execution_name):
    """
    Get the original input for an execution by constructing the ARN directly.
    This works for executions regardless of age.
    
    State machine ARN format: arn:aws:states:region:account:stateMachine:machine-name
    Execution ARN format: arn:aws:states:region:account:execution:machine-name:execution-name
    """
    try:
        # Extract the state machine name from the ARN
        state_machine_name = state_machine_arn.split(':')[-1]
        # Build execution ARN by replacing 'stateMachine' with 'execution' and appending execution name
        base_arn = ':'.join(state_machine_arn.split(':')[:-2])
        execution_arn = f"{base_arn}:execution:{state_machine_name}:{execution_name}"
        
        # Get execution details directly using the constructed ARN
        exec_details = sfn_client.describe_execution(executionArn=execution_arn)
        original_input = json.loads(exec_details.get('input', '{}'))
        return original_input
        
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        if error_code == 'ExecutionDoesNotExist':
            print(f"⚠️  Warning: Execution '{execution_name}' does not exist")
        else:
            print(f"⚠️  Warning: Could not get input for {execution_name}: {error_code}")
        return {}
    except json.JSONDecodeError as e:
        print(f"⚠️  Warning: Could not parse input JSON for {execution_name}: {e}")
        return {}


def retry_execution(sfn_client, state_machine_arn, execution_name, dry_run=False):
    """
    Create a retry execution for the given execution name.
    """
    try:
        # Find the base name and next retry count
        base_execution_name, retry_count = find_highest_retry_count(sfn_client, state_machine_arn, execution_name)
        
        # Generate new execution name using the base name
        new_execution_name = f"{base_execution_name}-retry-{retry_count}"
        
        # Get original input (try to find the original execution, not the retry)
        original_input = get_execution_input(sfn_client, state_machine_arn, base_execution_name)
        
        # If we can't find the base execution, try the provided execution name
        if not original_input:
            original_input = get_execution_input(sfn_client, state_machine_arn, execution_name)
        
        if dry_run:
            print(f"🔍 DRY RUN: Would create retry for {execution_name}")
            if base_execution_name != execution_name:
                print(f"   Base execution: {base_execution_name}")
            print(f"   New execution name: {new_execution_name}")
            print(f"   Retry count: {retry_count}")
            print(f"   Original input keys: {list(original_input.keys()) if original_input else 'None'}")
            return True, new_execution_name
    
        
        # Start the retry execution
        response = sfn_client.start_execution(
            stateMachineArn=state_machine_arn,
            name=new_execution_name,
            input=json.dumps(original_input) if original_input else '{}'
        )
        
        new_execution_arn = response.get('executionArn')
        print(f"✅ Successfully created retry for: {execution_name}")
        if base_execution_name != execution_name:
            print(f"   Base execution: {base_execution_name}")
        print(f"   New execution: {new_execution_name}")
        print(f"   New ARN: {new_execution_arn}")

        return True, new_execution_name
        
    except ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
        error_message = e.response.get('Error', {}).get('Message', str(e))
        
        print(f"❌ Failed to create retry for {execution_name}")
        print(f"   Error: {error_code} - {error_message}")
        
        if error_code == 'ExecutionAlreadyExists':
            print(f"   💡 Tip: An execution with name '{new_execution_name}' already exists")
        
        return False, None
    
    except Exception as e:
        print(f"❌ Unexpected error creating retry for {execution_name}: {e}")
        return False, None


def load_executions_from_file(filename):
    """
    Load execution names from a text file (one per line).
    Ignores empty lines and lines starting with #.
    """
    try:
        with open(filename, 'r') as f:
            executions = []
            for line in f:
                line = line.strip()
                # Skip empty lines and comments
                if line and not line.startswith('#'):
                    executions.append(line)
        return executions
    except FileNotFoundError:
        print(f"❌ File not found: {filename}")
        return None
    except Exception as e:
        print(f"❌ Error reading file {filename}: {e}")
        return None


def bulk_bulk_start_new_execution(sfn_client, state_machine_arn, execution_names, dry_run=False):
    """
    Create retry executions for multiple execution names.
    """
    if not execution_names:
        print("ℹ️  No executions to retry.")
        return
    
    print(f"\n🚀 {'DRY RUN: Would create retries for' if dry_run else 'Creating retries for'} {len(execution_names)} execution(s)")
    print("-" * 80)
    
    successful_retries = []
    failed_retries = []
    
    for i, execution_name in enumerate(execution_names, 1):
        print(f"\n{i}/{len(execution_names)}. Processing: {execution_name}")
        
        success, new_name = retry_execution(
            sfn_client, state_machine_arn, execution_name, dry_run
        )
        
        if success:
            successful_retries.append((execution_name, new_name))
        else:
            failed_retries.append(execution_name)
    
    # Summary
    print("\n" + "=" * 80)
    print("📊 RETRY SUMMARY")
    print("=" * 80)
    
    if dry_run:
        print(f"🔍 DRY RUN COMPLETE")
        print(f"✅ Would successfully create: {len(successful_retries)} retry execution(s)")
        print(f"❌ Would fail to create: {len(failed_retries)} retry execution(s)")
        print(f"\n💡 To actually create the retry executions, run the command again without --dry-run")
    else:
        print(f"✅ Successfully created: {len(successful_retries)} retry execution(s)")
        print(f"❌ Failed to create: {len(failed_retries)} retry execution(s)")
    
    if successful_retries:
        print(f"\n✅ Successful retries:")
        for original, new in successful_retries:
            print(f"   {original} → {new}")
    
    if failed_retries:
        print(f"\n❌ Failed retries:")
        for failed in failed_retries:
            print(f"   {failed}")
    
    return len(successful_retries), len(failed_retries)


def main():
    """
    Main function to handle command line arguments and orchestrate the retry process.
    """
    print("🔄 AWS Step Functions Execution Retry Tool")
    print("=" * 50)
    
    # Set up command line argument parsing
    parser = argparse.ArgumentParser(
        description='Create retry executions for specified execution names with incremented retry counts',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python bulk_start_new_execution.py MyStateMachine --executions ProcessData-001 ProcessData-002
  python bulk_start_new_execution.py MyStateMachine --executions ProcessData-001-retry-1 --dry-run
  python bulk_start_new_execution.py MyStateMachine --file executions.txt --dry-run
  python bulk_start_new_execution.py MyStateMachine --executions ProcessData-001 --dry-run

Execution Name Format:
  Input: ProcessData-001
  Output: ProcessData-001-retry-0 (first retry)
  Output: ProcessData-001-retry-1 (second retry, if retry-0 already exists)
  
  Input: ProcessData-001-retry-2 (already a retry)
  Output: ProcessData-001-retry-3 (next retry in sequence)

File Format (for --file option):
  One execution name per line:
  ProcessData-001
  ProcessData-002-retry-1
  ProcessData-003
        """
    )
    
    parser.add_argument('state_machine_name', 
                       help='Name of the Step Functions state machine')
    
    # Mutually exclusive group for input methods
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('--executions', '-e',
                           nargs='+',
                           help='List of execution names to retry')
    input_group.add_argument('--file', '-f',
                           help='File containing execution names (one per line)')
    
    parser.add_argument('--dry-run',
                       action='store_true',
                       help='Preview what would be created without actually starting new executions')
    
    args = parser.parse_args()
    
    print(f"🎯 Target state machine: {args.state_machine_name}")
    print(f"🔍 Mode: {'DRY RUN' if args.dry_run else 'LIVE EXECUTION'}")
    
    # Get execution names from either command line or file
    if args.executions:
        execution_names = args.executions
        print(f"📝 Input method: Command line ({len(execution_names)} execution(s))")
    else:
        execution_names = load_executions_from_file(args.file)
        if execution_names is None:
            sys.exit(1)
        print(f"📝 Input method: File '{args.file}' ({len(execution_names)} execution(s))")
    
    print(f"📋 Executions to retry: {', '.join(execution_names[:5])}")
    if len(execution_names) > 5:
        print(f"    ... and {len(execution_names) - 5} more")
    
    if args.dry_run:
        print("⚠️  DRY RUN MODE: No actual executions will be started")
    else:
        print("🚨 LIVE MODE: New retry executions will be started!")
        
        # Confirm before proceeding in live mode
        response = input("\nDo you want to proceed with creating retry executions? (y/N): ").strip().lower()
        if response != 'y':
            print("❌ Aborted by user")
            sys.exit(0)
    
    # Setup AWS session
    sfn_client = setup_aws_session()
    
    # Find state machine ARN
    state_machine_arn = find_state_machine_arn(sfn_client, args.state_machine_name)
    if not state_machine_arn:
        sys.exit(1)
    
    print(f"✓ Found state machine ARN: {state_machine_arn}")
    
    # Perform bulk retry
    successful_count, failed_count = bulk_bulk_start_new_execution(
        sfn_client, state_machine_arn, execution_names, dry_run=args.dry_run
    )
    
    # Final status
    if not args.dry_run:
        total_attempted = successful_count + failed_count
        print(f"\n🏁 FINAL RESULT: {successful_count}/{total_attempted} retry executions successfully created")
        
        if successful_count > 0:
            print("\n💡 Next steps:")
            print("- Monitor the new retry executions in the AWS Step Functions console")
            print("- Check CloudWatch logs for any issues")
            print("- Verify that the retry executions complete successfully")


if __name__ == "__main__":
    main()