import boto3
import argparse
import time
import base64
from botocore.exceptions import ClientError
from user_data import user_data

# Stałe
PREFIX = 's464953'
REGION = 'us-east-1'
AMI_ID = 'ami-0b5eea76982371e91'
INSTANCE_TYPE = 't2.micro'
MIN_INSTANCES = 2
MAX_INSTANCES = 5
LOAD_BALANCER_PORT = 80
TARGET_PORT = 8080

def clean_resources(ec2, prefix, default_vpc):
    # Usuń security group
    try:
        security_groups = ec2.meta.client.describe_security_groups(
            Filters=[{'Name': 'group-name', 'Values': [f'{prefix}-group']}]
        )['SecurityGroups']
        for sg in security_groups:
            print(f"Usuwam Security Group: {sg['GroupId']}")
            ec2.meta.client.delete_security_group(GroupId=sg['GroupId'])
    except ClientError as e:
        if e.response['Error']['Code'] != 'InvalidGroup.NotFound':
            raise

    # Usuń key pair
    try:
        print(f"Usuwam Key Pair: {prefix}-key")
        ec2.meta.client.delete_key_pair(KeyName=f'{prefix}-key')
    except ClientError as e:
        if e.response['Error']['Code'] != 'InvalidKeyPair.NotFound':
            raise

    # Usuń Launch Template
    try:
        launch_templates = ec2.meta.client.describe_launch_templates(
            Filters=[{'Name': 'launch-template-name', 'Values': [f'{prefix}-LaunchTemplate']}]
        )['LaunchTemplates']
        for lt in launch_templates:
            print(f"Usuwam Launch Template: {lt['LaunchTemplateId']}")
            ec2.meta.client.delete_launch_template(LaunchTemplateId=lt['LaunchTemplateId'])
    except ClientError as e:
        if e.response['Error']['Code'] != 'InvalidLaunchTemplateName.NotFound':
            raise

    # Usuń Auto Scaling Group
    autoscaling = boto3.client('autoscaling', region_name=REGION)
    try:
        autoscaling.delete_auto_scaling_group(
            AutoScalingGroupName=f'{prefix}-AutoScalingGroup',
            ForceDelete=True
        )
        print(f"Usuwam Auto Scaling Group: {prefix}-AutoScalingGroup")
    except ClientError as e:
        if e.response['Error']['Code'] != 'ValidationError':
            raise

    # Usuń Load Balancer i Target Group
    elbv2 = boto3.client('elbv2', region_name=REGION)
    try:
        target_groups = elbv2.describe_target_groups(
            Names=[f'{prefix}-TargetGroup']
        )['TargetGroups']
        for tg in target_groups:
            print(f"Usuwam Target Group: {tg['TargetGroupArn']}")
            elbv2.delete_target_group(TargetGroupArn=tg['TargetGroupArn'])
    except ClientError as e:
        if e.response['Error']['Code'] != 'TargetGroupNotFound':
            raise

    try:
        load_balancers = elbv2.describe_load_balancers(
            Names=[f'{prefix}-LoadBalancer']
        )['LoadBalancers']
        for lb in load_balancers:
            print(f"Usuwam Load Balancer: {lb['LoadBalancerArn']}")
            elbv2.delete_load_balancer(LoadBalancerArn=lb['LoadBalancerArn'])
    except ClientError as e:
        if e.response['Error']['Code'] != 'LoadBalancerNotFound':
            raise


def main(aws_access_key_id, aws_secret_access_key, aws_session_token, default_vpc):
    ec2 = boto3.resource(
        'ec2',
        region_name=REGION,
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key,
        aws_session_token=aws_session_token,
    )

    print("Czyszczenie istniejących zasobów...")
    clean_resources(ec2, PREFIX, default_vpc)

    print("Tworzenie zasobów...")
    # Tworzenie key pair
    key_pair = ec2.create_key_pair(
        KeyName=PREFIX + '-key',
        KeyType='ed25519',
        KeyFormat='pem',
    )

    # Tworzenie security group
    security_group = ec2.create_security_group(
        Description=PREFIX + '-group',
        GroupName=PREFIX + '-group',
        VpcId=default_vpc,
    )

    security_group.authorize_ingress(
        GroupId=security_group.group_id,
        IpPermissions=[
            {
                'IpProtocol': 'tcp',
                'FromPort': LOAD_BALANCER_PORT,
                'ToPort': LOAD_BALANCER_PORT,
                'IpRanges': [{'CidrIp': '0.0.0.0/0'}],
            },
        ]
    )

    # Tworzenie Launch Template
    ec2_client = boto3.client('ec2', region_name=REGION)
    launch_template = ec2_client.create_launch_template(
        LaunchTemplateName=PREFIX + '-LaunchTemplate',
        LaunchTemplateData={
            'ImageId': AMI_ID,
            'InstanceType': INSTANCE_TYPE,
            'KeyName': key_pair.name,
            'UserData': base64.b64encode(user_data.encode('utf-8')).decode('utf-8'),
            'SecurityGroupIds': [security_group.group_id],
        }
    )

    elbv2 = boto3.client('elbv2', region_name=REGION)
    target_group = elbv2.create_target_group(
        Name=PREFIX + '-TargetGroup',
        Protocol='HTTP',
        Port=TARGET_PORT,
        VpcId=default_vpc,
        TargetType='instance',
        HealthCheckProtocol='HTTP',
        HealthCheckPort=str(TARGET_PORT),
        HealthCheckPath='/factors/2',
    )
    target_group_arn = target_group['TargetGroups'][0]['TargetGroupArn']

    # Tworzenie Load Balancer
    load_balancer = elbv2.create_load_balancer(
        Name=PREFIX + '-LoadBalancer',
        Subnets=[default_vpc],  # Użyj odpowiednich subnetów
        Scheme='internet-facing',
        Type='application',
        IpAddressType='ipv4',
    )
    load_balancer_arn = load_balancer['LoadBalancers'][0]['LoadBalancerArn']

    elbv2.create_listener(
        LoadBalancerArn=load_balancer_arn,
        Protocol='HTTP',
        Port=LOAD_BALANCER_PORT,
        DefaultActions=[
            {
                'Type': 'forward',
                'TargetGroupArn': target_group_arn,
            },
        ]
    )

    # Tworzenie Auto Scaling Group
    autoscaling = boto3.client('autoscaling', region_name=REGION)
    autoscaling.create_auto_scaling_group(
        AutoScalingGroupName=PREFIX + '-AutoScalingGroup',
        LaunchTemplate={
            'LaunchTemplateId': launch_template['LaunchTemplate']['LaunchTemplateId'],
            'Version': '$Latest',
        },
        MinSize=MIN_INSTANCES,
        MaxSize=MAX_INSTANCES,
        DesiredCapacity=MIN_INSTANCES,
        VPCZoneIdentifier=default_vpc,  # Użyj odpowiednich subnetów
        TargetGroupARNs=[target_group_arn],
    )

    # Tworzenie reguł skalowania
    cloudwatch = boto3.client('cloudwatch', region_name=REGION)

    # Alarm skalowania w górę
    cloudwatch.put_metric_alarm(
        AlarmName=PREFIX + '-ScaleUp',
        MetricName='RequestCountPerTarget',
        Namespace='AWS/ApplicationELB',
        Statistic='Sum',
        Period=60,
        Threshold=100,
        ComparisonOperator='GreaterThanThreshold',
        EvaluationPeriods=2,
        AlarmActions=[
            autoscaling.put_scaling_policy(
                AutoScalingGroupName=PREFIX + '-AutoScalingGroup',
                PolicyName=PREFIX + '-ScaleUpPolicy',
                ScalingAdjustment=1,
                AdjustmentType='ChangeInCapacity',
            )['PolicyARN']
        ],
        Dimensions=[
            {'Name': 'TargetGroup', 'Value': target_group_arn},
        ],
    )

    # Alarm skalowania w dół
    cloudwatch.put_metric_alarm(
        AlarmName=PREFIX + '-ScaleDown',
        MetricName='RequestCountPerTarget',
        Namespace='AWS/ApplicationELB',
        Statistic='Sum',
        Period=60,
        Threshold=10,
        ComparisonOperator='LessThanThreshold',
        EvaluationPeriods=2,
        AlarmActions=[
            autoscaling.put_scaling_policy(
                AutoScalingGroupName=PREFIX + '-AutoScalingGroup',
                PolicyName=PREFIX + '-ScaleDownPolicy',
                ScalingAdjustment=-1,
                AdjustmentType='ChangeInCapacity',
            )['PolicyARN']
        ],
        Dimensions=[
            {'Name': 'TargetGroup', 'Value': target_group_arn},
        ],
    )

    print(f'Load Balancer URL: {load_balancer["LoadBalancers"][0]["DNSName"]}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Create AWS infrastructure")
    parser.add_argument('--aws_access_key_id', required=True, help='AWS Access Key ID')
    parser.add_argument('--aws_secret_access_key', required=True, help='AWS Secret Access Key')
    parser.add_argument('--aws_session_token', required=True, help='AWS Session Token')
    parser.add_argument('--default_vpc', required=True, help='Default VPC ID')

    args = parser.parse_args()
    main(args.aws_access_key_id, args.aws_secret_access_key, args.aws_session_token, args.default_vpc)