import boto3
from botocore.exceptions import ClientError
import time

# A script to create an Application Load Balancer (ALB) with a target group
# using Boto3.

# --- Configuration ---
REGION = "us-east-1"
ALB_NAME = "MyBoto3Alb"
SG_NAME = "MyBoto3AlbSG"
TG_NAME = "MyBoto3AlbTG"

ec2_client = boto3.client('ec2', region_name=REGION)
elbv2_client = boto3.client('elbv2', region_name=REGION)

def get_default_vpc_and_subnets():
    """Gets the default VPC ID and two default subnet IDs."""
    print("--- Getting Default VPC and Subnet IDs ---")
    try:
        vpcs = ec2_client.describe_vpcs(Filters=[{'Name': 'is-default', 'Values': ['true']}])['Vpcs']
        if not vpcs:
            raise Exception("No default VPC found.")
        vpc_id = vpcs[0]['VpcId']
        print(f"Default VPC ID: {vpc_id}")

        subnets = ec2_client.describe_subnets(Filters=[{'Name': 'vpc-id', 'Values': [vpc_id]}, {'Name': 'default-for-az', 'Values': ['true']}])['Subnets']
        if len(subnets) < 2:
            raise Exception("Not enough default subnets found in the default VPC.")
        subnet_ids = [s['SubnetId'] for s in subnets[:2]] # Take first two
        print(f"Default Subnet IDs: {subnet_ids}")
        return vpc_id, subnet_ids
    except ClientError as e:
        print(f"Error getting VPC/Subnet info: {e}")
        raise

def create_security_group(vpc_id):
    """Creates a Security Group for ALB."""
    print(f"\n--- Creating Security Group: {SG_NAME} ---")
    try:
        sg_response = ec2_client.create_security_group(
            GroupName=SG_NAME,
            Description="Allow HTTP traffic for ALB",
            VpcId=vpc_id
        )
        sg_id = sg_response['GroupId']
        print(f"Security Group '{SG_NAME}' created with ID: {sg_id}")

        ec2_client.authorize_security_group_ingress(
            GroupId=sg_id,
            IpPermissions=[
                {'IpProtocol': 'tcp', 'FromPort': 80, 'ToPort': 80, 'IpRanges': [{'CidrIp': '0.0.0.0/0'}]}
            ]
        )
        print("Inbound HTTP rule added to Security Group.")
        return sg_id
    except ClientError as e:
        if e.response['Error']['Code'] == 'InvalidGroup.Duplicate':
            print(f"Security Group '{SG_NAME}' already exists. Fetching ID.")
            response = ec2_client.describe_security_groups(GroupNames=[SECURITY_GROUP_NAME])
            return response['SecurityGroups'][0]['GroupId']
        else:
            print(f"Error creating security group: {e}")
            raise

def create_alb(subnet_ids, sg_id):
    """Creates an Application Load Balancer."""
    print(f"\n--- Creating Application Load Balancer: {ALB_NAME} ---")
    try:
        alb_response = elbv2_client.create_load_balancer(
            Name=ALB_NAME,
            Subnets=subnet_ids,
            SecurityGroups=[sg_id],
            Scheme='internet-facing',
            Type='application'
        )
        alb_arn = alb_response['LoadBalancers'][0]['LoadBalancerArn']
        print(f"ALB created with ARN: {alb_arn}. Waiting for it to be active...")
        elbv2_client.get_waiter('load_balancer_available').wait(LoadBalancerArns=[alb_arn])
        print("ALB is active.")
        return alb_arn
    except ClientError as e:
        print(f"Error creating ALB: {e}")
        raise

def create_target_group(vpc_id):
    """Creates an HTTP target group."""
    print(f"\n--- Creating Target Group: {TG_NAME} ---")
    try:
        tg_response = elbv2_client.create_target_group(
            Name=TG_NAME,
            Protocol='HTTP',
            Port=80,
            VpcId=vpc_id,
            HealthCheckProtocol='HTTP',
            HealthCheckPath='/',
            HealthCheckIntervalSeconds=30,
            HealthCheckTimeoutSeconds=5,
            HealthyThresholdCount=2,
            UnhealthyThresholdCount=2
        )
        tg_arn = tg_response['TargetGroups'][0]['TargetGroupArn']
        print(f"Target Group created with ARN: {tg_arn}")
        return tg_arn
    except ClientError as e:
        print(f"Error creating Target Group: {e}")
        raise

def create_listener(alb_arn, tg_arn):
    """Creates an HTTP listener for the ALB."""
    print(f"\n--- Creating Listener for ALB ---")
    try:
        elbv2_client.create_listener(
            LoadBalancerArn=alb_arn,
            Protocol='HTTP',
            Port=80,
            DefaultActions=[{'Type': 'forward', 'TargetGroupArn': tg_arn}]
        )
        print("Listener created.")
    except ClientError as e:
        print(f"Error creating Listener: {e}")
        raise

def cleanup_resources(alb_arn, tg_arn, sg_id):
    """Cleans up all created resources."""
    print(f"\n--- Cleaning up resources ---")

    # Delete ALB
    if alb_arn:
        print(f"Deleting ALB '{ALB_NAME}'...")
        try:
            elbv2_client.delete_load_balancer(LoadBalancerArn=alb_arn)
            elbv2_client.get_waiter('load_balancer_not_exists').wait(LoadBalancerArns=[alb_arn])
            print("ALB deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'LoadBalancerNotFound':
                print(f"ALB '{ALB_NAME}' not found, skipping deletion.")
            else:
                print(f"Error deleting ALB: {e}")

    # Delete Target Group
    if tg_arn:
        print(f"Deleting Target Group '{TG_NAME}'...")
        try:
            elbv2_client.delete_target_group(TargetGroupArn=tg_arn)
            print("Target Group deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'TargetGroupNotFound':
                print(f"Target Group '{TG_NAME}' not found, skipping deletion.")
            else:
                print(f"Error deleting Target Group: {e}")

    # Delete Security Group
    if sg_id:
        print(f"Deleting Security Group '{SG_NAME}'...")
        try:
            ec2_client.delete_security_group(GroupId=sg_id)
            print("Security Group deleted.")
        except ClientError as e:
            if e.response['Error']['Code'] == 'InvalidGroup.NotFound':
                print(f"Security Group '{SG_NAME}' not found, skipping deletion.")
            elif e.response['Error']['Code'] == 'DependencyViolation':
                print(f"Security Group '{SG_NAME}' is still in use. Retrying deletion after a short delay.")
                time.sleep(10)
                ec2_client.delete_security_group(GroupId=sg_id)
                print("Security Group deleted.")
            else:
                print(f"Error deleting security group: {e}")

def main():
    alb_arn = None
    tg_arn = None
    sg_id = None
    try:
        vpc_id, subnet_ids = get_default_vpc_and_subnets()
        sg_id = create_security_group(vpc_id)
        alb_arn = create_alb(subnet_ids, sg_id)
        tg_arn = create_target_group(vpc_id)
        create_listener(alb_arn, tg_arn)

        alb_description = elbv2_client.describe_load_balancers(LoadBalancerArns=[alb_arn])
        alb_dns_name = alb_description['LoadBalancers'][0]['DNSName']

        print("\n--- ALB Setup Complete! ---")
        print(f"ALB DNS Name: {alb_dns_name}")
        print(f"You can now register targets (e.g., EC2 instances) to the target group '{TG_NAME}'.")

        input("Press Enter to delete the ALB and clean up resources...")

    except ClientError as e:
        print(f"An AWS client error occurred: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
    finally:
        cleanup_resources(alb_arn, tg_arn, sg_id)
        print("\n--- ALB demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()