import boto3
import logging
import os

def lambda_handler(event, context):
    ECS_CLUSTER = os.getenv('ECS_CLUSTER')
    TASK_DEFINITION = os.getenv('TASK_DEFINITION')
    TASK_REVISION = os.getenv('TASK_REVISION')
    PUBLIC_IP = os.getenv('PUBLIC_IP', 'ENABLED')
    MAX_CONCURRENCY = int(os.getenv('MAX_CONCURRENCY', '1'))
    # List of used subnets and security groups.
    # Env variable should be comma-separated list for both.
    SUBNETS = os.getenv('SUBNETS', '').lower().replace(" ", "").split(',')
    SECURITY_GROUPS = os.getenv('SECURITY_GROUPS', '').lower().replace(" ", "").split(',')

    LOGLEVEL = os.getenv('LOGLEVEL', 'INFO').upper()
    logger = logging.getLogger('')
    logger.setLevel(LOGLEVEL)

    if ("ECS_CLUSTER" not in os.environ
      or "TASK_DEFINITION" not in os.environ
      or "TASK_REVISION" not in os.environ
      or "SUBNETS" not in os.environ):
        raise ValueError("Some of environment variables haven't been set.")

    ecs_client = boto3.client('ecs')

    # Check running tasks
    task_list = ecs_client.list_tasks(
        cluster=ECS_CLUSTER,
        family=f'{TASK_DEFINITION}',
        desiredStatus='RUNNING'
    )
    if len(task_list['taskArns']) < MAX_CONCURRENCY:
        response = ecs_client.run_task(
            cluster=ECS_CLUSTER,
            taskDefinition=f'{TASK_DEFINITION}:{TASK_REVISION}',
            count=1,
            launchType='FARGATE',
            networkConfiguration={
                'awsvpcConfiguration': {
                    'subnets': SUBNETS ,
                    'securityGroups': SECURITY_GROUPS ,
                    'assignPublicIp': PUBLIC_IP ,
                },
            },
            propagateTags='TASK_DEFINITION'
        )
        logger.info(response)
    else:
        logger.info(f"Too many tasks running. Currently running tasks - {len(task_list['taskArns'])}")