#!/bin/bash
# Function to edit the DaemonSet with kubectl patch
patch_daemonset() {
    FILE="aws-node-daemon-set.json"
    kubectl -n kube-system get daemonset aws-node -o json > $FILE
    jq '.spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms
      |= map(
          .matchExpressions |= map(
            select(
              .key != "eks.amazonaws.com/nodegroup" 
              or .operator != "NotIn" 
              or (.values | length != 1 or .[0] != "x-compute")
            )
          )
        )
      | .spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[0].matchExpressions += [{
          "key": "eks.amazonaws.com/nodegroup",
          "operator": "NotIn",
          "values": ["x-compute"]
        }]
    ' $FILE > $FILE.tmp && mv $FILE.tmp $FILE
    kubectl apply -f $FILE
    rm $FILE
}

# Function to export and modify the new DaemonSet
edit_daemonset() {
    FILE="exo-aws-node-daemon-set.json"
    # Export the current configuration to a JSON file
    kubectl -n kube-system get daemonset aws-node -o json > $FILE
    # Get rid of some metadata info
    jq 'del(.metadata.resourceVersion, .metadata.uid)' "$FILE" > "$FILE.tmp" && mv "$FILE.tmp" "$FILE"
    # Change nodeAffinity to schedule on x-compute nodes
    jq '(.spec.template.spec.affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms[].matchExpressions[]) 
        |= if .key == "eks.amazonaws.com/nodegroup" 
            and .operator == "NotIn" 
            and (.values | length > 0) 
            and .values[0] == "x-compute" 
            then .operator = "In" 
            else . 
        end
    ' $FILE > "$FILE.tmp" && mv "$FILE.tmp" $FILE
    # Update names
    sed -i 's/\"name\": \"aws-node\"/\"name\": \"exo-aws-node\"/g' $FILE
    sed -i 's/\"k8s-app\": \"aws-node\"/\"k8s-app\": \"exo-aws-node\"/g' $FILE
    sed -i 's/\"app.kubernetes.io\/name\": \"aws-node\"/\"app.kubernetes.io\/name\": \"exo-aws-node\"/g' $FILE
    sed -i 's/\"app.kubernetes.io\/instance\": \"aws-vpc-cni\"/\"app.kubernetes.io\/instance\": \"exo-aws-vpc-cni\"/g' $FILE
    # Replace image for aws-vpc-cni-driver
    jq '(.spec.template.spec.containers[].image) 
        |= if test("amazon-k8s-cni:") 
            then "public.ecr.aws/u8h5n6o4/aws-vpc-cni-driver:v1.19.0" 
            else . 
        end
    ' $FILE > $FILE.tmp && mv $FILE.tmp $FILE
    # Replace image for amazon-k8s-cni-init
    jq '(.spec.template.spec.initContainers[].image) 
        |= if test("amazon-k8s-cni-init:") 
            then "public.ecr.aws/u8h5n6o4/aws-vpc-cni-init:v1.19.0" 
            else . 
        end
    ' "$FILE" > "$FILE.tmp" && mv "$FILE.tmp" "$FILE"
    # Replace --enable-network-policy flag from true to false
    sed -i 's/--enable-network-policy=true/--enable-network-policy=false/g' $FILE
    # Apply the updated configuration to the cluster
    kubectl apply -f $FILE
    rm $FILE
}

# Function to set environment variables for the new DaemonSet
tune_daemonset() {
    kubectl -n kube-system set env daemonset exo-aws-node \
      AWS_VPC_K8S_CNI_EXTERNALSNAT=true \
      ENABLE_PREFIX_DELEGATION=true \
      MAX_ENI=1 \
      WARM_PREFIX_TARGET=13 \
      WARM_IP_TARGET=0 \
      MINIMUM_IP_TARGET=0 \
      DISABLE_NETWORK_RESOURCE_PROVISIONING=true \
      POD_SECURITY_GROUP_ENFORCING_MODE=standard
}
# Main function to execute the steps
main() {
    echo "Patching DaemonSet configuration..."
    patch_daemonset
    echo "Editing DaemonSet configuration..."
    edit_daemonset
    echo "Tuning DaemonSet environment variables..."
    tune_daemonset
    echo "DaemonSet configuration and tuning completed."
}
# Execute the main function
main