PROFILE=
POLICY=arn:aws:iam::aws:policy/CloudWatchAgentServerPolicy
USE_IAM_SERVICE_ACCOUNTS=
LOG_RETENTION_DAYS=7

if [ -z "$REGION" ]
then
    region_arg=""
else
    region_arg="--region $REGION"
fi

if [ -z "$PROFILE" ]
then
    profile_arg=""
else
    profile_arg="--profile $PROFILE"
fi

aws_version=$(aws --version | cut -d " " -f1 | cut -d/ -f2)
aws_major_version=$(echo $aws_version | cut -d. -f1)
aws_minor_version=$(echo $aws_version | cut -d. -f2)
if [ $aws_major_version -lt 2 ] || [ $aws_major_version -eq 2 ] && [ $aws_minor_version -lt 8 ]
then
    echo "WARNING: Outdated AWS CLI version detected - $aws_version, the script might fail."
    echo "Follow these instructions to update your AWS CLI version: https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html#getting-started-install-instructions"
fi

if [ -z "$USE_IAM_SERVICE_ACCOUNTS" ]
then
    node_roles=()

    echo "Fetch node IAM roles from EKS Nodegroups"
    node_groups=$(aws eks list-nodegroups --cluster-name $CLUSTER $region_arg $profile_arg --query "nodegroups" --output text)
    for node_group in ${node_groups[@]}; do
        node_role=$(aws eks describe-nodegroup --cluster-name $CLUSTER $region_arg $profile_arg --nodegroup-name $node_group --query "nodegroup.nodeRole" --output text)
        node_roles+=($node_role)
    done

    echo "Fetch node IAM instance roles from EC2 instances"
    ec2_instance_profile_ids=($(aws ec2 describe-instances $region_arg $profile_arg --filter Name=tag-key,Values=kubernetes.io/cluster/$CLUSTER --query "Reservations[*].Instances[*].IamInstanceProfile.Id" --output text | sed 's/\s/\n/g' | sed -r '/^\s*$/d' | sort -u))
    for profile_id in ${ec2_instance_profile_ids[@]}; do
        echo "Fetch node IAM roles for EC2 instance profile with ID $profile_id"
        instance_roles=$(aws iam list-instance-profiles $profile_arg --query "InstanceProfiles[?InstanceProfileId=='$profile_id'].Roles[*].Arn" --output text)
        node_roles+=(${instance_roles[@]})
    done

    echo "Fetch node IAM instance profiles from auto-scaling groups"
    lt_instance_profile_names=()
    launch_template_ids=($(aws autoscaling describe-auto-scaling-groups $region_arg $profile_arg --filter Name=tag-key,Values=kubernetes.io/cluster/$CLUSTER --query "AutoScalingGroups[*].MixedInstancesPolicy.LaunchTemplate.LaunchTemplateSpecification.LaunchTemplateId" --output text | sed 's/\s/\n/g' | sed -r '/^\s*$/d' | sort -u))
    launch_template_ids+=($(aws autoscaling describe-auto-scaling-groups $region_arg $profile_arg --filter Name=tag-key,Values=kubernetes.io/cluster/$CLUSTER --query "AutoScalingGroups[*].LaunchTemplate.LaunchTemplateId" --output text | sed 's/\s/\n/g' | sed -r '/^\s*$/d' | sort -u))
    for launch_template_id in ${launch_template_ids[@]}; do
        lt_instance_profile_names+=($(aws ec2 describe-launch-template-versions $region_arg $profile_arg --launch-template-id ${launch_template_id} --versions $Default $Latest --query "LaunchTemplateVersions[*].LaunchTemplateData.IamInstanceProfile.Name" --output text | sed 's/\s/\n/g' | sort -u))
    done
    for profile_name in ${lt_instance_profile_names[@]}; do
        echo "Fetch node IAM roles for EC2 instance profile named $profile_name"
        instance_roles=$(aws iam list-instance-profiles $profile_arg --query "InstanceProfiles[?InstanceProfileName=='$profile_name'].Roles[*].Arn" --output text)
        node_roles+=(${instance_roles[@]})
    done

    unique_node_roles=($(for role in "${node_roles[@]}"; do echo "${role}"; done | sort -u))
    for node_role in ${unique_node_roles[@]}; do
        role_name="${node_role#*/}"
        echo "Attach IAM policy $POLICY to role $role_name"
        aws iam attach-role-policy \
            --policy-arn $POLICY \
            --role-name $role_name \
            $profile_arg
    done
else
    echo "Check if the cluster has an ODIC provider."
    oidc_id=$(aws eks describe-cluster --name $CLUSTER $region_arg $profile_arg --query "cluster.identity.oidc.issuer" --output text | cut -d '/' -f 5)
    
    if [[ $(aws iam list-open-id-connect-providers | grep $oidc_id) ]]
    then
        echo "You already have an IAM OIDC provider associated with your cluster."
    else
        echo "No IAM OIDC identity provider found. Creating one for your cluster."
        eksctl utils associate-iam-oidc-provider --cluster $CLUSTER --approve
    fi

    eksctl create iamserviceaccount \
        --name aws-otel-sa \
        --namespace aws-otel-eks \
        --cluster $CLUSTER \
        --role-name "$CLUSTER-aws-otel-sa" \
        --attach-policy-arn $POLICY \
        --approve
fi

if [ ! -z "$LOG_RETENTION_DAYS" ]
then
    log_group_name="/aws/containerinsights/$CLUSTER/performance"
    log_group_created=$(aws logs describe-log-groups --log-group-name-prefix $log_group_name $region_arg $profile_arg --query 'length(logGroups[])' --output text)
    if [[ $log_group_created  -eq 0 ]]
    then
        echo "Create log group $log_group_name"
        aws logs create-log-group --log-group-name $log_group_name $region_arg $profile_arg
    fi

    echo "Set log group $log_group_name retention to $LOG_RETENTION_DAYS days"
    aws logs put-retention-policy --log-group-name $log_group_name --retention-in-days $LOG_RETENTION_DAYS $region_arg $profile_arg
fi

ami_ids=($(aws ec2 describe-instances --filters "Name=tag-key,Values=kubernetes.io/cluster/$CLUSTER" $region_arg $profile_arg --query 'Reservations[*].Instances[*].ImageId' --output text | sort -u))
for launch_template_id in ${launch_template_ids[@]}; do
    ami_ids+=($(aws ec2 describe-launch-template-versions $region_arg $profile_arg --launch-template-id ${launch_template_id} --versions $Default $Latest --query "LaunchTemplateVersions[*].LaunchTemplateData.ImageId" --output text | sed 's/\s/\n/g' | sort -u))
done

ami_names=($(aws ec2 describe-images $region_arg $profile_arg --image-ids ${ami_ids[@]} --query 'Images[*].Name' --output text))
echo "Detected AMIs: ${ami_names[@]}"

pattern="^bottlerocket-aws-k8s-1\.2[0-9]"
bottlerocket_amis_count=0
other_amis_count=0
for ami_name in ${ami_names[@]}; do
    if [[ $ami_name =~ $pattern ]]; then
        ((bottlerocket_amis_count++))
    else
        ((other_amis_count++))
    fi
done

echo "Update kubeconfig"
aws eks update-kubeconfig --name $CLUSTER $region_arg $profile_arg

if [[ $bottlerocket_amis_count -gt 0 && $other_amis_count -eq 0 ]]; then
    echo "Bottlerocket AMI detected, installing Amazon OTEL Agent with dockershim.sock patch."
    echo "Bottlerocket overrides containerd's socket to be dockershim.sock: https://github.com/bottlerocket-os/bottlerocket/issues/1126#issuecomment-698045504"
    echo "Issue description: https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/ContainerInsights-troubleshooting.html#ContainerInsights-troubleshooting-bottlerocket"
    curl -o otel-container-insights-infra.yaml https://uniskai-eu-templates.s3.amazonaws.com/eks/otel-container-insights-infra.yaml
    sed -i 's/path: \/run\/containerd\/containerd.sock/path: \/run\/dockershim.sock/g' otel-container-insights-infra.yaml

    echo "Applying Amazon OTEL Agent manifest with /run/dockershim.sock patch"
    kubectl apply -f otel-container-insights-infra.yaml

elif [[ $bottlerocket_amis_count -gt 0 && $other_amis_count -gt 0 ]]; then
    echo "Mix of BottleRocket and other AMI types detected."
    echo "If you ever face an issue where billing data appears, but grouping by workload properties (Service, Controller, etc.) does not work,"
    echo "and the cluster view is empty; please try to install Amazon OTEL Agent with dockershim.sock patch"
    echo "Details: https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/ContainerInsights-troubleshooting.html#ContainerInsights-troubleshooting-bottlerocket"
    echo "Installing Amazon OTEL Agent"
    curl https://uniskai-eu-templates.s3.amazonaws.com/eks/otel-container-insights-infra.yaml | kubectl apply -f -

else
    echo "Installing Amazon OTEL Agent"
    curl https://uniskai-eu-templates.s3.amazonaws.com/eks/otel-container-insights-infra.yaml | kubectl apply -f -
fi