import boto3
import argparse
import datetime

def cleanup_unused_load_balancers(
    region_name='us-east-1',
    delete_unused=False,
    idle_threshold_hours=24 * 7 # 7 days
):
    """
    Identifies and optionally deletes unused Application Load Balancers (ALBs)
    and Network Load Balancers (NLBs) that are incurring costs.

    Args:
        region_name (str): The AWS region to scan.
        delete_unused (bool): If True, deletes identified unused load balancers.
        idle_threshold_hours (int): Load balancers with no processed bytes for this duration
                                    will be considered unused.
    """
    elbv2_client = boto3.client('elbv2', region_name=region_name)
    cloudwatch_client = boto3.client('cloudwatch', region_name=region_name)

    print(f"Starting unused Load Balancer cleanup in region {region_name}...")
    if not delete_unused:
        print("*** DRY RUN MODE: No changes will be applied. ***")

    unused_lbs = []

    # Helper function to check if an LB has been idle
    def is_load_balancer_idle(lb_arn, lb_type):
        metric_name = 'ProcessedBytes'
        namespace = 'AWS/ApplicationELB' if lb_type == 'application' else 'AWS/NetworkELB'

        end_time = datetime.datetime.now(datetime.timezone.utc)
        start_time = end_time - datetime.timedelta(hours=idle_threshold_hours)

        response = cloudwatch_client.get_metric_statistics(
            Namespace=namespace,
            MetricName=metric_name,
            Dimensions=[
                {'Name': 'LoadBalancer', 'Value': lb_arn.split('/')[3] + '/' + lb_arn.split('/')[4] + '/' + lb_arn.split('/')[5]}
            ],
            StartTime=start_time,
            EndTime=end_time,
            Period=3600, # 1 hour
            Statistics=['Sum']
        )

        total_processed_bytes = 0
        for dp in response['Datapoints']:
            total_processed_bytes += dp['Sum']
        
        return total_processed_bytes == 0

    try:
        # 1. Describe Application Load Balancers (ALBs)
        print("\n>>> Step 1: Describing Application Load Balancers (ALBs)...")
        alb_response = elbv2_client.describe_load_balancers(LoadBalancerType='application')
        albs = alb_response['LoadBalancers']
        print(f"   Found {len(albs)} ALBs.")

        for alb in albs:
            alb_arn = alb['LoadBalancerArn']
            alb_name = alb['LoadBalancerName']
            print(f"   Checking ALB: {alb_name} ({alb_arn})")
            if is_load_balancer_idle(alb_arn, 'application'):
                unused_lbs.append({'Name': alb_name, 'Arn': alb_arn, 'Type': 'ALB', 'Reason': f'No processed bytes for {idle_threshold_hours} hours.'})
                print(f"      ALB '{alb_name}' identified as unused.")
            else:
                print(f"      ALB '{alb_name}' is active.")

        # 2. Describe Network Load Balancers (NLBs)
        print("\n>>> Step 2: Describing Network Load Balancers (NLBs)...")
        nlb_response = elbv2_client.describe_load_balancers(LoadBalancerType='network')
        nlbs = nlb_response['LoadBalancers']
        print(f"   Found {len(nlbs)} NLBs.")

        for nlb in nlbs:
            nlb_arn = nlb['LoadBalancerArn']
            nlb_name = nlb['LoadBalancerName']
            print(f"   Checking NLB: {nlb_name} ({nlb_arn})")
            if is_load_balancer_idle(nlb_arn, 'network'):
                unused_lbs.append({'Name': nlb_name, 'Arn': nlb_arn, 'Type': 'NLB', 'Reason': f'No processed bytes for {idle_threshold_hours} hours.'})
                print(f"      NLB '{nlb_name}' identified as unused.")
            else:
                print(f"      NLB '{nlb_name}' is active.")

    except Exception as e:
        print(f"Error describing load balancers: {e}")
        return

    # 3. Generate Report and Optionally Delete
    print("\n>>> Step 3: Generating report and performing cleanup (if enabled)...")
    report_output = []
    if unused_lbs:
        report_output.append("--- Unused Load Balancers Found ---")
        for lb_info in unused_lbs:
            report_output.append(f"  Type: {lb_info['Type']}")
            report_output.append(f"  Name: {lb_info['Name']}")
            report_output.append(f"  ARN: {lb_info['Arn']}")
            report_output.append(f"  Reason: {lb_info['Reason']}")
            report_output.append("----------------------------------------")

            if delete_unused:
                print(f"   Deleting {lb_info['Type']} '{lb_info['Name']}'...")
                try:
                    elbv2_client.delete_load_balancer(LoadBalancerArn=lb_info['Arn'])
                    print(f"      Successfully initiated deletion for {lb_info['Name']}.")
                except Exception as e:
                    print(f"      Error deleting {lb_info['Type']} '{lb_info['Name']}': {e}")
    else:
        report_output.append("No unused load balancers found based on criteria.")

    if report_file_path:
        with open(report_file_path, 'w') as f:
            for line in report_output:
                f.write(line + '\n')
        print(f"Report saved to '{report_file_path}'.")
    else:
        for line in report_output:
            print(line)

    print("\nUnused Load Balancer cleanup completed.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Identify and optionally delete unused Load Balancers.")
    parser.add_argument("--region", default="us-east-1", help="AWS region to scan (default: us-east-1).")
    parser.add_argument("--delete-unused", action="store_true", help="If set, deletes identified unused load balancers.")
    parser.add_argument("--idle-threshold-hours", type=int, default=168, help="Load balancers with no processed bytes for this duration (hours) will be considered unused (default: 168 hours / 7 days).")
    parser.add_argument("--report-file-path", help="Optional. Path to a file to save the report. If not provided, prints to console.")

    args = parser.parse_args()

    cleanup_unused_load_balancers(
        region_name=args.region,
        delete_unused=args.delete_unused,
        idle_threshold_hours=args.idle_threshold_hours,
        report_file_path=args.report_file_path
    )