#!/bin/bash

# Set your variables
# Cluster name
CLUSTER_NAME="$1"

# Service name
SERVICE_NAME="$2"

# task definition family
TASK_DEFINITION_FAMILY="$3"

# Sensor name
SENSOR_NAME="$4"

# Sensor image in AWS repository
SENSOR_IMAGE="$5"

# Controller name or URL
CONTROLLER_NAME="$6"

# Controller port number
CONTROLLER_PORT="$7"

# Environment type - dev/prod
ENVIRONMENT_TYPE="$8"

# TAP interface name
TAP_INTERFACE="$9"

# TAP filter
TAP_FILTER="${10}"

# Relay protocol
RELAY_PROTOCOL="${11}"

# AWS region
AWS_REGION="${12}"

# Function to check if jq is installed
check_jq() {
    if ! command -v jq &> /dev/null; then
        echo "jq is required but not installed. Please install jq and try again."
        exit 1
    fi
}

# Function to get the current task definition
get_task_definition() {
    aws ecs describe-task-definition --task-definition $TASK_DEFINITION_FAMILY
    if [ $? -ne 0 ]; then
        echo "Failed to retrieve the task definitions, exiting"
        exit 1
    fi
}

# Function to add or update the sensor sidecar
add_or_update_sensor() {
    local task_def="$1"
    local taskDefinitionFamily=$(echo "$task_def" | jq -r '.taskDefinition.family')
    local awsgroup="/ecs/${taskDefinitionFamily}"
    local sensor_exists=$(echo "$task_def" | jq --arg SENSOR_NAME "$SENSOR_NAME" '.taskDefinition.containerDefinitions[] | select(.name == $SENSOR_NAME) | length')

    if [ -z "$sensor_exists" ] || [ "$sensor_exists" == "0" ]; then
        # Add new sensor sidecar while retaining all existing containers
        echo "$task_def" | jq --arg SENSOR_NAME "$SENSOR_NAME" --arg SENSOR_IMAGE "$SENSOR_IMAGE" --arg CONTROLLER_NAME "$CONTROLLER_NAME" --arg CONTROLLER_PORT "$CONTROLLER_PORT" --arg ENVIRONMENT_TYPE "$ENVIRONMENT_TYPE" --arg TAP_INTERFACE "$TAP_INTERFACE" --arg TAP_FILTER "$TAP_FILTER" --arg RELAY_PROTOCOL "$RELAY_PROTOCOL" --arg AWS_REGION "$AWS_REGION" --arg AWS_GROUP "$awsgroup" '.taskDefinition.containerDefinitions += [{
            "name": $SENSOR_NAME,
            "image": $SENSOR_IMAGE,
            "essential": true,
            "memory": 512,
            "cpu": 256,
            "environment": [
                {
                    "name": "REMOTE_CONTROLLER_SERVER_NAME",
                    "value": $CONTROLLER_NAME
                },
                {
                    "name": "REMOTE_CONTROLLER_SERVER_PORT",
                    "value": $CONTROLLER_PORT
                },
                {
                    "name": "ENVIRONMENT",
                    "value": $ENVIRONMENT_TYPE
                },
                {
                    "name": "TAP_INTERFACE",
                    "value": $TAP_INTERFACE
                },
                {
                    "name": "TAP_FILTER",
                    "value": $TAP_FILTER
                },
                {
                    "name": "RELAY_PROTOCOL",
                    "value": $RELAY_PROTOCOL
                },
                {
                    "name": "TAP_PROFILE",
                    "value": "low"
                },
                {
                    "name": "MAX_SAAS_QUEUE_HOLDING_CAP",
                    "value": "4096"
                },
                {
                    "name": "SAAS_CONNECTOR_LIMIT",
                    "value": "10"
                }
            ],
            "logConfiguration": {
                "logDriver": "awslogs",
                "options": {
                    "awslogs-group": $AWS_GROUP,
                    "awslogs-create-group": "true",
                    "awslogs-region": $AWS_REGION,
                    "awslogs-stream-prefix": "ecs"
                }
            }
        }]'
    else
        # Update existing sensor sidecar while retaining all other containers
        echo "$task_def" | jq --arg SENSOR_NAME "$SENSOR_NAME" --arg SENSOR_IMAGE "$SENSOR_IMAGE" '.taskDefinition.containerDefinitions = [.taskDefinition.containerDefinitions[] | if .name == $SENSOR_NAME then . + {"image": $SENSOR_IMAGE} else . end]'
    fi
}

# Function to register new task definition
register_task_definition() {
    local new_task_def="$1"
    echo "$new_task_def" | jq '{
        family: .taskDefinition.family,
        taskRoleArn: .taskDefinition.taskRoleArn,
        executionRoleArn: .taskDefinition.executionRoleArn,
        networkMode: (.taskDefinition.networkMode // null),
        containerDefinitions: .taskDefinition.containerDefinitions,
        volumes: .taskDefinition.volumes,
        placementConstraints: .taskDefinition.placementConstraints,
        requiresCompatibilities: .taskDefinition.requiresCompatibilities,
        cpu: .taskDefinition.cpu,
        memory: .taskDefinition.memory
    } | del(.networkMode | select(. == null))' > updated_task_def.json

    aws ecs register-task-definition --cli-input-json file://updated_task_def.json
    if [ $? -ne 0 ]; then
        echo "Failed to register task definitions"
        echo "$new_task_def"
        exit 1
    fi
}

# Function to update the ECS service
update_ecs_service() {
    local new_task_def_arn="$1"
    aws ecs update-service --cluster $CLUSTER_NAME --service $SERVICE_NAME --task-definition $new_task_def_arn
    if [ $? -ne 0 ]; then
        echo "Failed to updated the service, exiting"
        exit 1
    fi
}

# Main execution
set -e
check_jq

echo "Fetching current task definition..."
current_task_def=$(get_task_definition)
echo "$current_task_def"

echo "Adding or updating sensor sidecar..."
new_task_def=$(add_or_update_sensor "$current_task_def")
echo "$new_task_def"

echo "Registering new task definition..."
registered_task_def=$(register_task_definition "$new_task_def")

new_task_def_arn=$(echo $registered_task_def | jq -r '.taskDefinition.taskDefinitionArn')

echo "Updating ECS service..."
update_ecs_service "$new_task_def_arn"

echo "Sidecar deployment process completed."