import boto3
import argparse
import json
import time

def setup_ec2_power_management_schedule(
    instance_tag_key,
    instance_tag_value,
    stop_cron_expression,
    start_cron_expression,
    region_name='us-east-1'
):
    """
    Sets up automated power management for EC2 instances based on tags and a schedule.
    This involves creating a Lambda function and two EventBridge (CloudWatch Events)
    rules to stop and start instances at specified times.

    Args:
        instance_tag_key (str): The tag key to identify target EC2 instances (e.g., 'AutoStop').
        instance_tag_value (str): The tag value to identify target EC2 instances (e.g., 'true').
        stop_cron_expression (str): Cron expression for stopping instances (e.g., 'cron(0 19 ? * MON-FRI *)').
        start_cron_expression (str): Cron expression for starting instances (e.g., 'cron(0 8 ? * MON-FRI *)').
        region_name (str): The AWS region.
    """
    lambda_client = boto3.client('lambda', region_name=region_name)
    events_client = boto3.client('events', region_name=region_name)
    iam_client = boto3.client('iam', region_name=region_name)

    lambda_function_name = f"ec2-power-manager-{instance_tag_key}-{instance_tag_value}"
    stop_rule_name = f"ec2-stop-rule-{instance_tag_key}-{instance_tag_value}"
    start_rule_name = f"ec2-start-rule-{instance_tag_key}-{instance_tag_value}"

    print(f"Setting up EC2 power management for instances with tag {instance_tag_key}={instance_tag_value} in {region_name}...")

    # 1. Create IAM Role for Lambda Function
    print("\n>>> Step 1: Creating/Updating IAM Role for Lambda function...")
    lambda_role_name = f"lambda-ec2-power-manager-role-{instance_tag_key}-{instance_tag_value}"
    assume_role_policy_document = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {"Service": "lambda.amazonaws.com"},
                "Action": "sts:AssumeRole"
            }
        ]
    }
    try:
        create_role_response = iam_client.create_role(
            RoleName=lambda_role_name,
            AssumeRolePolicyDocument=json.dumps(assume_role_policy_document)
        )
        lambda_role_arn = create_role_response['Role']['Arn']
        print(f"   Created IAM Role: {lambda_role_arn}")
    except iam_client.exceptions.EntityAlreadyExistsException:
        lambda_role_arn = iam_client.get_role(RoleName=lambda_role_name)['Role']['Arn']
        print(f"   IAM Role '{lambda_role_name}' already exists: {lambda_role_arn}")
    except Exception as e:
        print(f"Error creating/getting IAM role: {e}")
        return

    # Attach policies to the role
    # AWSLambdaBasicExecutionRole for CloudWatch Logs access
    iam_client.attach_role_policy(
        RoleName=lambda_role_name,
        PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
    )
    # Policy for EC2 start/stop actions
    ec2_policy_document = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": [
                    "ec2:StartInstances",
                    "ec2:StopInstances",
                    "ec2:DescribeInstances"
                ],
                "Resource": "*" # Granular resource specification is recommended in production
            }
        ]
    }
    ec2_policy_name = f"ec2-power-manager-policy-{instance_tag_key}-{instance_tag_value}"
    try:
        iam_client.put_role_policy(
            RoleName=lambda_role_name,
            PolicyName=ec2_policy_name,
            PolicyDocument=json.dumps(ec2_policy_document)
        )
        print(f"   Attached EC2 Start/Stop policy to role '{lambda_role_name}'.")
    except Exception as e:
        print(f"Error attaching EC2 policy: {e}")
        return

    # Give IAM time to propagate
    time.sleep(10)

    # 2. Create/Update Lambda Function
    print("\n>>> Step 2: Creating/Updating Lambda Function...")
    lambda_code = f"""
import boto3
import os

def lambda_handler(event, context):
    ec2 = boto3.client('ec2', region_name=os.environ['AWS_REGION'])
    action = event['action']
    tag_key = os.environ['INSTANCE_TAG_KEY']
    tag_value = os.environ['INSTANCE_TAG_VALUE']

    print(f"Received event to {action} EC2 instances with tag {tag_key}={tag_value}")

    filters = [
        {{'Name': f'tag:{tag_key}', 'Values': [tag_value]}},
        {{'Name': 'instance-state-name', 'Values': ['running', 'stopped']}}
    ]

    instances = ec2.describe_instances(Filters=filters)['Reservations']
    instance_ids = []
    for reservation in instances:
        for instance in reservation['Instances']:
            instance_ids.append(instance['InstanceId'])

    if not instance_ids:
        print("No instances found matching the criteria.")
        return

    if action == 'stop':
        print(f"Stopping instances: {instance_ids}")
        ec2.stop_instances(InstanceIds=instance_ids)
    elif action == 'start':
        print(f"Starting instances: {instance_ids}")
        ec2.start_instances(InstanceIds=instance_ids)
    else:
        print(f"Invalid action: {action}")

    return {{'statusCode': 200, 'body': f'{action} initiated for {len(instance_ids)} instances.'}}
"""

    try:
        # Check if function exists
        lambda_client.get_function(FunctionName=lambda_function_name)
        # Update existing function
        lambda_client.update_function_code(
            FunctionName=lambda_function_name,
            ZipFile=lambda_code.encode('utf-8')
        )
        lambda_client.update_function_configuration(
            FunctionName=lambda_function_name,
            Runtime='python3.9',
            Handler='index.lambda_handler',
            Role=lambda_role_arn,
            Environment={'Variables': {'INSTANCE_TAG_KEY': instance_tag_key, 'INSTANCE_TAG_VALUE': instance_tag_value, 'AWS_REGION': region_name}},
            Timeout=30,
            MemorySize=128
        )
        print(f"   Updated Lambda function: {lambda_function_name}")
    except lambda_client.exceptions.ResourceNotFoundException:
        # Create new function
        lambda_client.create_function(
            FunctionName=lambda_function_name,
            Runtime='python3.9',
            Handler='index.lambda_handler',
            Role=lambda_role_arn,
            Code={'ZipFile': lambda_code.encode('utf-8')},
            Environment={'Variables': {'INSTANCE_TAG_KEY': instance_tag_key, 'INSTANCE_TAG_VALUE': instance_tag_value, 'AWS_REGION': region_name}},
            Timeout=30,
            MemorySize=128
        )
        print(f"   Created Lambda function: {lambda_function_name}")
    except Exception as e:
        print(f"Error creating/updating Lambda function: {e}")
        return

    # 3. Create/Update EventBridge Rules
    print("\n>>> Step 3: Creating/Updating EventBridge Rules...")

    # Stop Rule
    try:
        events_client.put_rule(
            Name=stop_rule_name,
            ScheduleExpression=stop_cron_expression,
            State='ENABLED',
            Description=f"Stop EC2 instances with tag {instance_tag_key}={instance_tag_value}"
        )
        events_client.put_targets(
            Rule=stop_rule_name,
            Targets=[
                {
                    'Id': '1',
                    'Arn': lambda_client.get_function(FunctionName=lambda_function_name)['Configuration']['FunctionArn'],
                    'Input': json.dumps({'action': 'stop'})
                },
            ]
        )
        # Add permission for EventBridge to invoke Lambda
        lambda_client.add_permission(
            FunctionName=lambda_function_name,
            StatementId=f"AllowExecutionFromEventBridge-{stop_rule_name}",
            Action='lambda:InvokeFunction',
            Principal='events.amazonaws.com',
            SourceArn=events_client.describe_rule(Name=stop_rule_name)['Arn']
        )
        print(f"   Created/Updated Stop Rule: {stop_rule_name}")
    except Exception as e:
        print(f"Error setting up Stop Rule: {e}")
        return

    # Start Rule
    try:
        events_client.put_rule(
            Name=start_rule_name,
            ScheduleExpression=start_cron_expression,
            State='ENABLED',
            Description=f"Start EC2 instances with tag {instance_tag_key}={instance_tag_value}"
        )
        events_client.put_targets(
            Rule=start_rule_name,
            Targets=[
                {
                    'Id': '1',
                    'Arn': lambda_client.get_function(FunctionName=lambda_function_name)['Configuration']['FunctionArn'],
                    'Input': json.dumps({'action': 'start'})
                },
            ]
        )
        # Add permission for EventBridge to invoke Lambda
        lambda_client.add_permission(
            FunctionName=lambda_function_name,
            StatementId=f"AllowExecutionFromEventBridge-{start_rule_name}",
            Action='lambda:InvokeFunction',
            Principal='events.amazonaws.com',
            SourceArn=events_client.describe_rule(Name=start_rule_name)['Arn']
        )
        print(f"   Created/Updated Start Rule: {start_rule_name}")
    except Exception as e:
        print(f"Error setting up Start Rule: {e}")
        return

    print("\nEC2 power management schedule setup completed.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Setup automated power management for EC2 instances.")
    parser.add_argument("--instance-tag-key", required=True, help="Tag key to identify target EC2 instances (e.g., 'AutoStop').")
    parser.add_argument("--instance-tag-value", required=True, help="Tag value to identify target EC2 instances (e.g., 'true').")
    parser.add_argument("--stop-cron-expression", required=True, help="Cron expression for stopping instances (e.g., 'cron(0 19 ? * MON-FRI *)').")
    parser.add_argument("--start-cron-expression", required=True, help="Cron expression for starting instances (e.g., 'cron(0 8 ? * MON-FRI *)').")
    parser.add_argument("--region", default="us-east-1", help="AWS region (default: us-east-1).")

    args = parser.parse_args()

    setup_ec2_power_management_schedule(
        instance_tag_key=args.instance_tag_key,
        instance_tag_value=args.instance_tag_value,
        stop_cron_expression=args.stop_cron_expression,
        start_cron_expression=args.start_cron_expression,
        region_name=args.region
    )