import boto3
import base64
from USER_DATA import user_data
import json
import time

# Configuration variables
PREFIX = "s464863"
AMI_ID = "ami-0b5eea76982371e91"
INSTANCE_TYPE = "t2.micro"
SECURITY_GROUP_NAME = f"{PREFIX}-sg"
REGION = "us-east-1"
AWS_RESOURCES_FILE="aws_resources.json"

# Object to store AWS resources
aws_resources = {}

def create_vpc(ec2_client):
    # Create VPC
    response = ec2_client.create_vpc(
        CidrBlock='10.0.0.0/16',
        TagSpecifications=[
            {
                'ResourceType': 'vpc',
                'Tags': [{'Key': 'Name', 'Value': f"{PREFIX}-vpc"}]
            }
        ]
    )
    vpc_id = response['Vpc']['VpcId']
    print(f"Created VPC with ID: {vpc_id}")

    # Enable DNS support and hostnames
    ec2_client.modify_vpc_attribute(VpcId=vpc_id, EnableDnsSupport={'Value': True})
    ec2_client.modify_vpc_attribute(VpcId=vpc_id, EnableDnsHostnames={'Value': True})
    print("Enabled DNS support and hostnames.")

    return vpc_id


def create_subnet(ec2_client, vpc_id, cidr_block, availability_zone):
    # Create subnet
    response = ec2_client.create_subnet(
        VpcId=vpc_id,
        CidrBlock=cidr_block,
        AvailabilityZone=availability_zone,
        TagSpecifications=[
            {
                'ResourceType': 'subnet',
                'Tags': [{'Key': 'Name', 'Value': f"{PREFIX}-subnet"}]
            }
        ]
    )
    subnet_id = response['Subnet']['SubnetId']
    print(f"Created Subnet with ID: {subnet_id}")

    # Public IP addresses for instances in the subnet
    ec2_client.modify_subnet_attribute(SubnetId=subnet_id, MapPublicIpOnLaunch={'Value': True})
    print("Configured subnet to auto-assign public IP addresses.")

    return subnet_id


def create_internet_gateway(ec2_client, vpc_id):
    # Create Internet Gateway
    response = ec2_client.create_internet_gateway(
        TagSpecifications=[
            {
                'ResourceType': 'internet-gateway',
                'Tags': [{'Key': 'Name', 'Value': f"{PREFIX}-igw"}]
            }
        ]
    )
    igw_id = response['InternetGateway']['InternetGatewayId']
    print(f"Created Internet Gateway with ID: {igw_id}")

    # Attach Internet Gateway to VPC
    ec2_client.attach_internet_gateway(InternetGatewayId=igw_id, VpcId=vpc_id)
    print("Attached Internet Gateway to VPC.")

    return igw_id


def create_route_table(ec2_client, vpc_id, subnet_id, igw_id):
    # Create Route Table
    response = ec2_client.create_route_table(
        VpcId=vpc_id,
        TagSpecifications=[
            {
                'ResourceType': 'route-table',
                'Tags': [{'Key': 'Name', 'Value': f"{PREFIX}-rt"}]
            }
        ]
    )
    route_table_id = response['RouteTable']['RouteTableId']
    print(f"Created Route Table with ID: {route_table_id}")

    # Create route to Internet Gateway
    ec2_client.create_route(
        RouteTableId=route_table_id,
        DestinationCidrBlock='0.0.0.0/0',
        GatewayId=igw_id
    )
    print("Added route to Internet Gateway in Route Table.")

    # Associate Route Table with Subnet
    ec2_client.associate_route_table(RouteTableId=route_table_id, SubnetId=subnet_id)
    print("Associated Route Table with Subnet.")

    return route_table_id


def create_key_pair(ec2_client, key_name, save_to_file):
    # Create key pair
    response = ec2_client.create_key_pair(KeyName=key_name)

    # Save private key to file
    private_key = response['KeyMaterial']
    with open(save_to_file, 'w') as file:
        file.write(private_key)

    print(f"Key pair '{key_name}' created and saved to '{save_to_file}'.")
    return response['KeyName']


def create_security_group(ec2_client, vpc_id):
    group_name = f"{PREFIX}-sg"

    # Create Security Group
    response = ec2_client.create_security_group(
        GroupName=group_name,
        Description=f"Security group for {PREFIX} webservices",
        VpcId=vpc_id
    )

    security_group_id = response['GroupId']
    print(f"Security group '{group_name}' created with ID: {security_group_id}")

    # Set ingress rules
    ec2_client.authorize_security_group_ingress(
        GroupId=security_group_id,
        IpPermissions=[
            {
                'IpProtocol': 'tcp',
                'FromPort': 22,
                'ToPort': 22,
                'IpRanges': [{'CidrIp': '0.0.0.0/0', 'Description': 'SSH access'}]
            },
            {
                'IpProtocol': 'tcp',
                'FromPort': 8080,
                'ToPort': 8080,
                'IpRanges': [{'CidrIp': '0.0.0.0/0', 'Description': 'HTTP access'}]
            }
        ]
    )

    print("Ingress rules added for SSH (22) and HTTP (80).")
    return security_group_id


def create_launch_template(ec2_client, key_name, security_group_id):
    # Create Launch Template
    response = ec2_client.create_launch_template(
        LaunchTemplateName=f"{PREFIX}-lt",
        LaunchTemplateData={
            'ImageId': AMI_ID,
            'InstanceType': INSTANCE_TYPE,
            'SecurityGroupIds': [security_group_id],
            'KeyName': key_name,
            'UserData': base64.b64encode(user_data.encode('utf-8')).decode('utf-8')
        }
    )

    print(f"Launch Template created with ID: {response['LaunchTemplate']['LaunchTemplateId']}")
    return response['LaunchTemplate']['LaunchTemplateId']


def create_target_group(elbv2_client, vpc_id):
    # Create Target Group
    response = elbv2_client.create_target_group(
        Name=f"{PREFIX}-tg",
        Protocol='HTTP',
        Port=8080,
        VpcId=vpc_id,
        TargetType='instance',
        IpAddressType='ipv4',
        HealthCheckProtocol='HTTP',
        HealthCheckPort='8080',
        HealthCheckPath='/factors/6',
        HealthCheckIntervalSeconds=30,
        HealthCheckTimeoutSeconds=5,
        HealthyThresholdCount=3,
        UnhealthyThresholdCount=3,
    )

    print(f"Target Group created with ARN: {response['TargetGroups'][0]['TargetGroupArn']}")
    return response['TargetGroups'][0]['TargetGroupArn']


def create_load_balancer(elbv2_client, subnet_ids, security_group_id):
    # Create Load Balancer
    response = elbv2_client.create_load_balancer(
        Name=f"{PREFIX}-lb",
        Subnets=subnet_ids,
        SecurityGroups=[security_group_id],
        Scheme='internet-facing',
        Type='application',
        IpAddressType='ipv4'
    )

    load_balancer_arn = response['LoadBalancers'][0]['LoadBalancerArn']
    print(f"Load Balancer created with ARN: {load_balancer_arn}")
    return load_balancer_arn


def create_load_balancer_listener(elbv2_client, load_balancer_arn, target_group_arn):
    # Create Load Balancer Listener
    response = elbv2_client.create_listener(
        LoadBalancerArn=load_balancer_arn,
        Protocol='HTTP',
        Port=8080,
        DefaultActions=[
            {
                'Type': 'forward',
                'TargetGroupArn': target_group_arn
            }
        ]
    )

    print(f"Listener created with ARN: {response['Listeners'][0]['ListenerArn']}")
    return response['Listeners'][0]['ListenerArn']


def create_auto_scaling_group(autoscaling_client, launch_template_id, target_group_arn, subnet_id_1, subnet_id_2):
    # Create Auto Scaling Group
    autoscaling_client.create_auto_scaling_group(
        AutoScalingGroupName=f"{PREFIX}-asg",
        LaunchTemplate={
            'LaunchTemplateId': launch_template_id,
            'Version': '$Latest'
        },
        MinSize=2,
        MaxSize=5,
        DesiredCapacity=2,
        TargetGroupARNs=[target_group_arn],
        AvailabilityZones=[f"{REGION}a", f"{REGION}b"],
        HealthCheckType='EC2',
        HealthCheckGracePeriod=90,
        VPCZoneIdentifier=f"{subnet_id_1},{subnet_id_2}",
        Tags=[
            {
                'Key': 'Name',
                'Value': f"{PREFIX}-asg",
                'PropagateAtLaunch': True
            }
        ]
    )

    describe_response = autoscaling_client.describe_auto_scaling_groups(
        AutoScalingGroupNames=[f"{PREFIX}-asg"]
    )

    asg_arn = describe_response['AutoScalingGroups'][0]['AutoScalingGroupARN']
    print(f"Auto Scaling Group created with ARN: {asg_arn}")
    return asg_arn

def put_scaling_policies(autoscaling_client, asg_arn):
    # Put Scaling Policy Up
    up = autoscaling_client.put_scaling_policy(
        AutoScalingGroupName=f"{PREFIX}-asg",
        PolicyName=f"{PREFIX}-scale-policy-up",
        PolicyType='SimpleScaling',
        AdjustmentType='ChangeInCapacity',
        ScalingAdjustment=1,
        Cooldown=60
    )

    print(f"Scale policy created with ARN: {up['PolicyARN']}")

    # Put Scaling Policy Down
    down = autoscaling_client.put_scaling_policy(
        AutoScalingGroupName=f"{PREFIX}-asg",
        PolicyName=f"{PREFIX}-scale-policy-down",
        PolicyType='SimpleScaling',
        AdjustmentType='ChangeInCapacity',
        ScalingAdjustment=-1,
        Cooldown=60
    )

    print(f"Scale policy created with ARN: {down['PolicyARN']}")
    return up['PolicyARN'], down['PolicyARN']

def create_cloudwatch_alarms(cloudwatch_client, sp_up, sp_down):
    # Create CloudWatch Alarms
    cloudwatch_client.put_metric_alarm(
        AlarmName=f"{PREFIX}-scale-alarm",
        ComparisonOperator='GreaterThanThreshold',
        EvaluationPeriods=1,
        MetricName='RequestCountPerTarget',
        Namespace='AWS/ApplicationELB',
        Period=60,
        Statistic='Sum',
        Threshold=50,
        ActionsEnabled=True,
        AlarmActions=[sp_up],
        AlarmDescription='Scale-up',
        Dimensions=[{
            'Name': 'AutoScalingGroupName',
            'Value': f"{PREFIX}-asg"
        }],
        Unit='Count'
    )

    print(f"Scale-up alarm created")

    # Scale-down alarm
    cloudwatch_client.put_metric_alarm(
        AlarmName=f"{PREFIX}-scale-down-alarm",
        ComparisonOperator='LessThanThreshold',
        EvaluationPeriods=1,
        MetricName='RequestCountPerTarget',
        Namespace='AWS/ApplicationELB',
        Period=60,
        Statistic='Average',
        Threshold=10,
        ActionsEnabled=True,
        AlarmActions=[sp_down],
        AlarmDescription='Scale-down alarm',
        Dimensions=[{
            'Name': 'AutoScalingGroupName',
            'Value': f"{PREFIX}-asg"
        }],
        Unit='Count'
    )

    print(f"Scale-down alarm created")

# Main function
def main():
    # EC2 client
    ec2_client = boto3.client(
        'ec2',
        region_name=REGION
    )

    # Load Balancer client
    elbv2_client = boto3.client(
        'elbv2',
        region_name=REGION
    )

    # Auto Scaling client
    autoscaling_client = boto3.client(
        'autoscaling',
        region_name=REGION
    )

    # CloudWatch client
    cloudwatch_client = boto3.client(
        'cloudwatch',
        region_name=REGION
    )

    # Create VPC
    print("Creating VPC...")
    vpc_id = create_vpc(ec2_client)

    # Create first subnet
    print("Creating first subnet...")
    subnet_id_1 = create_subnet(ec2_client, vpc_id, cidr_block='10.0.1.0/24', availability_zone='us-east-1a')

    # Create second subnet
    print("Creating second subnet...")
    subnet_id_2 = create_subnet(ec2_client, vpc_id, cidr_block='10.0.2.0/24', availability_zone='us-east-1b')

    # Create Internet Gateway
    print("Creating Internet Gateway...")
    igw_id = create_internet_gateway(ec2_client, vpc_id)

    # Create Route Table
    print("Creating Route Table...")
    route_table_id = create_route_table(ec2_client, vpc_id, subnet_id_1, igw_id)

    # Create key pair
    print("Creating key pair...")
    key_name = create_key_pair(ec2_client, f"{PREFIX}-key", f"{PREFIX}-key.pem")

    # Create security group
    print("Creating security group...")
    security_group_id = create_security_group(ec2_client, vpc_id)

    # Create Launch Template
    print("Creating Launch Template...")
    launch_template_id = create_launch_template(ec2_client, key_name, security_group_id)

    # Create Target Group for ALB
    print("Creating Target Group...")
    target_group_arn = create_target_group(elbv2_client, vpc_id)

    # Create Load Balancer
    print("Creating Load Balancer...")
    load_balancer_arn = create_load_balancer(elbv2_client, [subnet_id_1, subnet_id_2], security_group_id)

    # Create Load Balancer Listener
    print("Creating Load Balancer Listener...")
    listener_arn = create_load_balancer_listener(elbv2_client, load_balancer_arn, target_group_arn)

    # Auto Scaling Group
    print("Creating Auto Scaling Group...")
    asg_arn = create_auto_scaling_group(autoscaling_client, launch_template_id, target_group_arn, subnet_id_1, subnet_id_2)

    # Scaling Policies
    print("Creating Scaling Policies...")
    sp_up, sp_down = put_scaling_policies(autoscaling_client, asg_arn)

    # CloudWatch Alarms
    print("Creating CloudWatch Alarms...")
    create_cloudwatch_alarms(cloudwatch_client, sp_up, sp_down)

    # Load Balancer DNS name
    response = elbv2_client.describe_load_balancers(
        LoadBalancerArns=[load_balancer_arn]
    )
    dns_name = response['LoadBalancers'][0]['DNSName']
    print(f"Load Balancer DNS name: {dns_name}")

    # Save AWS resources to file
    aws_resources = {
        'vpc_id': vpc_id,
        'subnet_ids': [subnet_id_1, subnet_id_2],
        'igw_id': igw_id,
        'route_table_id': route_table_id,
        'key_name': key_name,
        'security_group_id': security_group_id,
        'launch_template_id': launch_template_id,
        'target_group_arn': target_group_arn,
        'load_balancer_arn': load_balancer_arn,
        'listener_arn': listener_arn,
        'asg_name': f"{PREFIX}-asg",
        'dns_name': dns_name,
    }

    with open(AWS_RESOURCES_FILE, 'w') as file:
        json.dump(aws_resources, file, indent=4)

if __name__ == "__main__":
    main()