import os
import boto3
import sys

hsm_region = os.getenv('AWS_HSM_REGION')
nlb_region = os.getenv('AWS_NLB_REGION')
nlb_name = os.getenv('NLB_NAME', 'hsm-nlb')

lb_client = boto3.client('elbv2', region_name=nlb_region)

try:
    existing_nlb = lb_client.describe_load_balancers(Names=[nlb_name])
    if existing_nlb['LoadBalancers']:
        # Loop through the NLBs and list their associated target groups
        for load_balancer in existing_nlb['LoadBalancers']:
            nlb_name = load_balancer['LoadBalancerName']
            nlb_arn = load_balancer['LoadBalancerArn']
            # List the target groups associated with the NLB
            target_groups = lb_client.describe_target_groups(LoadBalancerArn=nlb_arn)
            lb_client.delete_load_balancer(
                LoadBalancerArn=nlb_arn
            )
            print(f"Deleted nlb: {nlb_name}")
            for tg in target_groups['TargetGroups']:
                lb_client.delete_target_group(
                    TargetGroupArn=tg['TargetGroupArn']
                )
                print(f"Deleted target group: {tg['TargetGroupName']}")
except:
    print("NLB with name hsm-nlb doesn't exist.")
    sys.exit()

print("Done to remove nlb and target groups")