import boto3
import os 
import threading
import time
from actions_logging.app_logging import logger



def wait_for_service_with_health_checks(cluster_name, service_name):
    ecs_client = boto3.client('ecs', region_name='us-east-1')
    while True:
        time.sleep(5) # wait for new task to spin up
        task_id = ecs_client.list_tasks(cluster=cluster_name, serviceName=service_name)['taskArns'][0]
        response = ecs_client.describe_tasks(cluster=cluster_name, tasks=[task_id])
        containers = response['tasks'][0]['containers']
        container = next((c for c in containers if c['name'] == f"{service_name}-{env_name}"), None)
        status = container['healthStatus'] if container else None
        logger.info(f"Status of service is: {status}")
        logger.info(f"Status of task {task_id}: {status}")
        if status == "HEALTHY":
            break
        else:
            logger.info(f"Waiting for service: {service_name} to reach steady state...")
            time.sleep(10)
        
        logger.info(f"Status of service: {service_name} Updated and stable")


def wait_for_service_without_health_checks(cluster_name, service_name):
    ecs_client = boto3.client('ecs', region_name='us-east-1')
    time.sleep(5) # wait for new task to spin up
    while True:
        response = ecs_client.describe_services(cluster=cluster_name, services=[service_name])
        status = response['services'][0]['deployments'][0]['rolloutState']
        logger.info(f"Status of service is: {status}")
        if status == 'COMPLETED':
            break
        elif status == 'FAILED':
            raise Exception(f"service: {service_name}  update failed")
        else:
            logger.info(f"Waiting for service: {service_name}  to reach steady state...")
            time.sleep(10)
    
    logger.info(f"Status of service: {service_name}  is: Updated and stable")


def update_service(cluster_name, service_name):
    list_ms_do_not_have_health_check = ["cd-poc-1","saferx-backend-cli","saferx-backend","sx-core-dns-public","sx-core-dns-private","scim-adapter","pzero-redis"]
    
    ecs_client = boto3.client('ecs', region_name='us-east-1')
    try:
        logger.info(f"Updating service: {service_name} ")
        ecs_client.update_service(cluster=cluster_name, service=service_name, forceNewDeployment=True)
        #time.sleep(2)  # Simulate update process
        
        if service_name in list_ms_do_not_have_health_check:
            wait_for_service_without_health_checks(cluster_name, service_name)
        else:
            wait_for_service_with_health_checks(cluster_name, service_name)
    except Exception as e:
        logger.error(f"An error occurred while updating service: {service_name}  {str(e)}")

def update_services(cluster_name):
    try:
        ecs_client = boto3.client('ecs', region_name='us-east-1')
        next_token = None
        service_arns = []  
        
        while True:
            if next_token:
                response = ecs_client.list_services(cluster=cluster_name, nextToken=next_token)
            else:
                response = ecs_client.list_services(cluster=cluster_name)
            next_token = response.get('nextToken')
            service_arns.extend(response.get('serviceArns', []))
            if not next_token:
                break
                
        threads = []
        for service_arn in service_arns:
            service_name = service_arn.split('/')[-1]
            if mss_to_igone:
                if service_name in svc_list_to_ignore:
                    logger.info(f"Skipping service: {service_name} ##################################################")
                    continue
            thread = threading.Thread(target=update_service, args=(cluster_name, service_name), name=service_name)
            threads.append(thread)
            thread.start()
            
        for thread in threads:
            thread.join()
            
    except Exception as e:
        logger.error(f"An error occurred: {str(e)}")


if __name__ == "__main__":
    env_name = os.environ.get('ENV_NAME')
    mss_to_igone = os.environ.get('MS_TO_IGNORE')
    if mss_to_igone:
        svc_list_to_ignore = mss_to_igone.split(',')
    cluster_name = f"{env_name}-micro-services"
    cluster_name_ec2 = f"{env_name}-ec2-micro-services"
    update_services(cluster_name)
    update_services(cluster_name_ec2)