import boto3
import json
import os
import time

ecs = boto3.client('ecs')
ecr = boto3.client('ecr')
code_deploy = boto3.client('codedeploy')
sts = boto3.client('sts')
sqs = boto3.client('sqs')

def lambda_handler(event, context):
    # Get the AWS account ID
    identity = sts.get_caller_identity()
    account_id = identity['Account']
    application_name = os.environ['APPLICATION_NAME']
    deployment_group_name = os.environ['DEPLOYMENT_GROUP_NAME']
    task_definition_family = os.environ['TASK_DEFINITION_FAMILY']
    deployment_queue_url = os.environ['QUEUE_URL']

    # Get the queue attributes to check the number of messages
    queue_attributes = sqs.get_queue_attributes(
        QueueUrl=deployment_queue_url,
        AttributeNames=['ApproximateNumberOfMessages']
    )

    number_of_messages = int(queue_attributes['Attributes']['ApproximateNumberOfMessages'])
    print(f"Number of messages in the queue: {number_of_messages}")

    for record in event['Records']:
        message_body = json.loads(record['body'])
        receipt_handle = record['receiptHandle']
        print(f"Message Body: {json.dumps(message_body)}")
        print(f"image-digest: {message_body['detail']['image-digest']}")

        try:
            describe_images_params = {
                'repositoryName': 'studio-api',
                'imageIds': [
                    {
                        'imageDigest': message_body['detail']['image-digest']
                    }
                ]
            }

            describe_images_response = ecr.describe_images(**describe_images_params)
            print(describe_images_response)
            print('Image details:', describe_images_response)

            image_tags = describe_images_response['imageDetails'][0]['imageTags']
            print('Image tags:', image_tags)

            # Extract the commitId from the imageTags
            commit_id = next((tag for tag in image_tags if tag != 'latest'), None)
            if not commit_id:
                raise Exception('No commitId found in image tags.')
            print('Commit ID:', commit_id)
            short_commit_id = commit_id[:7]
            print('Short Commit ID:', short_commit_id)
            # Fetch the latest task definition ARN
            task_def_response = ecs.list_task_definitions(
                familyPrefix=task_definition_family,
                sort='DESC',
                maxResults=1
            )

            if not application_name or not task_def_response['taskDefinitionArns']:
                raise Exception('No task definitions found for the specified family.')

            latest_task_definition_arn = task_def_response['taskDefinitionArns'][0]

            # Describe the latest task definition
            describe_task_def_response = ecs.describe_task_definition(
                taskDefinition=latest_task_definition_arn
            )

            container_definitions = describe_task_def_response['taskDefinition']['containerDefinitions']
            print('Container Definitions:', container_definitions)

            # Update the container definitions with the new image tag
            for container in container_definitions:
                if container['name'] == 'StudioApiContainer1':
                    image_parts = container['image'].split(':')
                    container['image'] = f"{image_parts[0]}:{commit_id}"

                    # Set DD_VERSION environment variable to the short commit id
                    dd_version_exists = False
                    for env_var in container['environment']:
                        if env_var['name'] == 'DD_VERSION':
                            env_var['value'] = short_commit_id
                            dd_version_exists = True
                            break
                    if not dd_version_exists:
                        container['environment'].append({
                            'name': 'DD_VERSION',
                            'value': short_commit_id
                        })
                        
                    print(f"Updated image for container {container['name']}: {container['image']}")

            # Copy all properties from the existing task definition
            new_task_definition = describe_task_def_response['taskDefinition'].copy()
            new_task_definition['containerDefinitions'] = container_definitions

            # Remove properties that should not be included in the new task definition
            del new_task_definition['status']
            del new_task_definition['taskDefinitionArn']
            del new_task_definition['requiresAttributes']
            del new_task_definition['compatibilities']
            del new_task_definition['registeredAt']
            del new_task_definition['registeredBy']
            del new_task_definition['revision']
            print(new_task_definition)

            # Register a new task definition with the updated container definitions
            new_task_def_response = ecs.register_task_definition(**new_task_definition)
            print(f"Registered new task definition: {new_task_def_response['taskDefinition']['taskDefinitionArn']}")

            # If the queue length is greater than 1, skip processing
            if number_of_messages > 1:
                print("Queue length is greater than 1, skipping processing.")
                continue
            # Create a new deployment
            deployment_response = code_deploy.create_deployment(
                applicationName=application_name,
                deploymentGroupName=deployment_group_name,
                revision={
                    'revisionType': 'AppSpecContent',
                    'appSpecContent': {
                        'content': json.dumps({
                            'version': 1,
                            'Resources': [
                                {
                                    'TargetService': {
                                        'Type': 'AWS::ECS::Service',
                                        'Properties': {
                                            'TaskDefinition': new_task_def_response['taskDefinition']['taskDefinitionArn'],
                                            'LoadBalancerInfo': {
                                                'ContainerName': 'StudioApiContainer1',
                                                'ContainerPort': 8005
                                            }
                                        }
                                    }
                                }
                            ],
                            'Hooks': [
                                {
                                    'BeforeInstall': 'DbMigrationFunctionBeforeInstall'
                                },
                                {
                                    'BeforeAllowTraffic': 'studio-api-pre-traffic-hook'
                                }
                            ]
                        })
                    }
                }
            )
            print(f"Created deployment: {deployment_response['deploymentId']}")

            # Poll the deployment status until it reaches a terminal state
            deployment_id = deployment_response['deploymentId']
            deployment_status = 'InProgress'
            while deployment_status == 'InProgress':
                time.sleep(10)  # Wait for 10 seconds before polling again

                deployment_status_response = code_deploy.get_deployment(deploymentId=deployment_id)
                deployment_status = deployment_status_response['deploymentInfo']['status']
                print(f"Current deployment status: {deployment_status}")

            if deployment_status == 'Succeeded':
                print('CodeDeploy deployment succeeded')
            else:
                print(f"CodeDeploy deployment failed with status: {deployment_status}")
                raise Exception(f"CodeDeploy deployment failed with status: {deployment_status}")

            # Delete the message from the queue after successful processing
            sqs.delete_message(
                QueueUrl=deployment_queue_url,
                ReceiptHandle=receipt_handle
            )
            print(f"Deleted message: {receipt_handle}")

        except Exception as e:
            print(f"Error processing message: {e}")
            raise e