import boto3
from botocore.exceptions import ClientError
import time

# A script to create an EFS file system and a mount target using Boto3.

# --- Configuration ---
REGION = "us-east-1"
EFS_FS_NAME = "MyBoto3EFSFileSystem"
SG_NAME = "MyBoto3EFS_SG"

ec2_client = boto3.client('ec2', region_name=REGION)
efs_client = boto3.client('efs', region_name=REGION)

def get_default_vpc_and_subnet():
    """Gets the default VPC ID and a default subnet ID."""
    print("--- Getting Default VPC and Subnet ID ---")
    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']
        vpc_cidr = vpcs[0]['CidrBlock']
        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 not subnets:
            raise Exception("No default subnet found in the default VPC.")
        subnet_id = subnets[0]['SubnetId']
        print(f"Default Subnet ID: {subnet_id}")
        return vpc_id, vpc_cidr, subnet_id
    except ClientError as e:
        print(f"Error getting VPC/Subnet info: {e}")
        raise

def create_security_group(vpc_id, vpc_cidr):
    """Creates a Security Group for EFS."""
    print(f"\n--- Creating Security Group: {SG_NAME} ---")
    try:
        sg_response = ec2_client.create_security_group(
            GroupName=SG_NAME,
            Description="Allow NFS traffic for EFS",
            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': 2049, 'ToPort': 2049, 'IpRanges': [{'CidrIp': vpc_cidr}]}
            ]
        )
        print("Inbound NFS 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=[SG_NAME])
            return response['SecurityGroups'][0]['GroupId']
        else:
            print(f"Error creating security group: {e}")
            raise

def create_efs_file_system():
    """Creates an EFS file system."""
    print(f"\n--- Creating EFS File System: {EFS_FS_NAME} ---")
    try:
        fs_response = efs_client.create_file_system(
            CreationToken=EFS_FS_NAME, # Unique string for idempotency
            PerformanceMode='generalPurpose',
            ThroughputMode='bursting',
            Encrypted=True,
            Tags=[{'Key': 'Name', 'Value': EFS_FS_NAME}]
        )
        fs_id = fs_response['FileSystemId']
        print(f"EFS File System created with ID: {fs_id}. Waiting for it to be available...")
        efs_client.get_waiter('file_system_available').wait(FileSystemId=fs_id)
        print("EFS File System is available.")
        return fs_id
    except ClientError as e:
        if e.response['Error']['Code'] == 'FileSystemAlreadyExists':
            print(f"EFS File System '{EFS_FS_NAME}' already exists. Fetching ID.")
            response = efs_client.describe_file_systems(CreationToken=EFS_FS_NAME)
            return response['FileSystems'][0]['FileSystemId']
        else:
            print(f"Error creating file system: {e}")
            raise

def create_mount_target(fs_id, subnet_id, sg_id):
    """Creates a mount target for the EFS file system."""
    print(f"\n--- Creating Mount Target for EFS File System '{fs_id}' ---")
    try:
        mt_response = efs_client.create_mount_target(
            FileSystemId=fs_id,
            SubnetId=subnet_id,
            SecurityGroups=[sg_id]
        )
        mt_id = mt_response['MountTargetId']
        print(f"Mount Target created with ID: {mt_id}. Waiting for it to be available...")
        efs_client.get_waiter('mount_target_available').wait(MountTargetId=mt_id)
        print("Mount Target is available.")
        return mt_id
    except ClientError as e:
        if e.response['Error']['Code'] == 'MountTargetConflict':
            print(f"Mount target already exists for file system '{fs_id}' in subnet '{subnet_id}'. Fetching ID.")
            response = efs_client.describe_mount_targets(FileSystemId=fs_id)
            for mt in response['MountTargets']:
                if mt['SubnetId'] == subnet_id:
                    return mt['MountTargetId']
            raise Exception("Existing mount target not found for specified subnet.")
        else:
            print(f"Error creating mount target: {e}")
            raise

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

    # Delete Mount Target
    if mt_id:
        print(f"Deleting Mount Target '{mt_id}'...")
        try:
            efs_client.delete_mount_target(MountTargetId=mt_id)
            # EFS doesn't have a 'mount_target_not_exists' waiter, so we poll
            while True:
                try:
                    efs_client.describe_mount_targets(MountTargetId=mt_id)
                    print("Mount target still exists, waiting...")
                    time.sleep(10)
                except ClientError as e:
                    if e.response['Error']['Code'] == 'MountTargetNotFound':
                        print("Mount Target deleted.")
                        break
                    else:
                        raise e
        except ClientError as e:
            print(f"Error deleting mount target: {e}")

    # Delete EFS File System
    if fs_id:
        print(f"Deleting EFS File System '{fs_id}'...")
        try:
            efs_client.delete_file_system(FileSystemId=fs_id)
            efs_client.get_waiter('file_system_deleted').wait(FileSystemId=fs_id)
            print("EFS File System deleted.")
        except ClientError as e:
            print(f"Error deleting file system: {e}")

    # Delete Security Group
    if sg_id:
        print(f"Deleting Security Group '{sg_id}'...")
        try:
            ec2_client.delete_security_group(GroupId=sg_id)
            print("Security Group deleted.")
        except ClientError as e:
            print(f"Error deleting security group: {e}")

def main():
    fs_id = None
    mt_id = None
    sg_id = None
    
    try:
        vpc_id, vpc_cidr, subnet_id = get_default_vpc_and_subnet()
        sg_id = create_security_group(vpc_id, vpc_cidr)
        fs_id = create_efs_file_system()
        mt_id = create_mount_target(fs_id, subnet_id, sg_id)

        print("\n--- EFS File System setup complete! ---")
        print(f"File System ID: {fs_id}")
        print(f"Mount Target ID: {mt_id}")
        print(f"You can now mount this EFS file system on an EC2 instance in subnet {subnet_id}.")

        input("Press Enter to delete the EFS 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(fs_id, mt_id, sg_id)
        print("\n--- EFS file system demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()