#!/usr/bin/env bash

set -xeuo pipefail

# When changing this script, remember to modify the corresponding github action
vendor=${1}
env=${2}
region=${3}
cluster=${4}
rsIpWhitelist=${5:-""}
dsIpWhitelist=${6:-""}

ROOT=$(cd $(dirname ${BASH_SOURCE[0]}) && pwd)

### BEGIN - Logic for existed region ###

update_ipset() {
  local region=$1
  local ipset_name=$2
  local add_ip_list=$3

  ipset_id=$(aws wafv2 list-ip-sets \
    --scope REGIONAL \
    --region ${region} \
    --query "IPSets[?Name=='${ipset_name}'].Id" \
    --output text)
  ipset=$(aws wafv2 get-ip-set \
    --scope REGIONAL \
    --region ${region} \
    --id ${ipset_id} \
    --name ${ipset_name})
  current_addresses=$(echo ${ipset} | jq -r '.IPSet.Addresses | join(" ")')
  add_ip_list=$(echo ${add_ip_list} | tr ',' ' ')
  new_addresses="${current_addresses} ${add_ip_list}"
  new_addresses=$(echo ${new_addresses} | tr ' ' '\n' | jq -R . | jq -s .)
  lock_token=$(echo ${ipset} | jq -r '.LockToken')
  aws wafv2 update-ip-set \
    --scope REGIONAL \
    --region ${region} \
    --id ${ipset_id} \
    --name ${ipset_name} \
    --addresses "${new_addresses}" \
    --lock-token ${lock_token}
}

update_azure_waf_ipset() {
  local region=$1
  local waf_name=$2
  local add_ip_list=$3
  local rule_name=$4

  echo "Updating Azure WAF IP set for $waf_name in $region"

  # Get the resource group name based on the cluster name
  resource_group="${cluster}-infra"

  # Get current WAF policy
  echo "Retrieving current WAF policy configuration"
  policy_id=$(az network application-gateway waf-policy show \
    --resource-group "${resource_group}" \
    --name "${waf_name}" \
    --query id -o tsv)

  if [ -z "$policy_id" ]; then
    echo "Error: WAF policy ${waf_name} not found in resource group ${resource_group}"
    return 1
  fi

  # Get current custom rules
  current_rules=$(az network application-gateway waf-policy custom-rule list \
    --resource-group "${resource_group}" \
    --policy-name "${waf_name}" -o json)

  rule_exists=$(echo "$current_rules" | jq -r --arg name "$rule_name" '.[] | select(.name==$name) | .name')

  # Process the IP list
  IFS=',' read -ra IPS <<< "${add_ip_list}"

  if [ -n "$rule_exists" ]; then
    echo "Updating existing IP whitelist rule"
    # Get current match conditions
    match_conditions=$(az network application-gateway waf-policy custom-rule show \
      --resource-group "${resource_group}" \
      --policy-name "${waf_name}" \
      --name "$rule_name" \
      --query matchConditions -o json)

    # Extract current IP addresses
    current_ips=$(echo "$match_conditions" | jq -r '.[0].matchValues[]')

    # Add new IPs to the list
    for ip in "${IPS[@]}"; do
      # Check if IP already exists
      ip_exists=$(echo "$current_ips" | grep -F "$ip" || true)
      if [ -z "$ip_exists" ]; then
        echo "Adding IP $ip to whitelist"
        match_conditions=$(echo "$match_conditions" | jq --arg ip "$ip" '.[0].matchValues += [$ip]')
      else
        echo "IP $ip already exists in whitelist"
      fi
    done

    # Nasty Hack: In order to make the Azure happy, we need to convert the JSON to the format that the CLI expects.
    # replace "matchValues" with "values", "matchVariables" with "variables", "negationConditon" with "negate"
    # and there is a f**king typo in negationConditon !!!
    match_conditions=$(echo "$match_conditions" | jq '.[0] |= with_entries(if .key == "matchValues" then .key = "values" elif .key == "matchVariables" then .key = "variables" elif .key == "negationConditon" then .key = "negate" else . end)')

    # Update the rule with new match conditions
    az network application-gateway waf-policy custom-rule update \
      --resource-group "${resource_group}" \
      --policy-name "${waf_name}" \
      --name "$rule_name" \
      --match-conditions "$match_conditions" \
      --priority 10 \
      --rule-type MatchRule \
      --action Allow
  else
    echo "custom rule not exists in waf $waf_name"
    exit 1
  fi
  echo "Successfully updated Azure WAF IP whitelist for $waf_name"
}

if [[ "${vendor}" == "azure" ]]; then
  if [ -n "${rsIpWhitelist}" ]; then
    rs_waf_name="${cluster}-rs-api"
    update_azure_waf_ipset ${region} ${rs_waf_name} ${rsIpWhitelist} "regionalwhitelist"
  fi
  exit 0
fi

if [ "${env}" == "dev" ] && [ "${region}" == "us-west-2" ]; then
  if [ -n "${dsIpWhitelist}" ]; then
    ds_ipset_name="centralized-storage-${env}-regional-${region}-ds-office"
    update_ipset ${region} ${ds_ipset_name} ${dsIpWhitelist}
  fi
  exit 0
fi

if [ "${env}" == "staging" ] && [ "${region}" == "us-west-2" ]; then
  if [ -n "${rsIpWhitelist}" ]; then
    rs_ipset_name="o11y-${env}-regional-${region}-rs-api"
    update_ipset ${region} ${rs_ipset_name} ${rsIpWhitelist}
  fi
  if [ -n "${dsIpWhitelist}" ]; then
    ds_ipset_name="o11y-${env}-regional-${region}-prometheus-api"
    update_ipset ${region} ${ds_ipset_name} ${dsIpWhitelist}
  fi
  exit 0
fi

if [ "${env}" == "prod" ]; then
  existed_regions=(
    "us-east-1"
    "us-east-2"
    "us-west-1"
    "us-west-2"
    "ap-northeast-1"
    "ap-northeast-2"
    "ap-southeast-1"
    "ap-southeast-3"
    "ap-south-1"
    "eu-central-1"
  )
  for existed_region in "${existed_regions[@]}"; do
    if [ "${region}" == "${existed_region}" ]; then
      if [ -n "${rsIpWhitelist}" ]; then
        rs_ipset_name="o11y-${env}-regional-${region}-rs-office"
        if [ "${region}" == "ap-southeast-3" ] || [ "${region}" == "ap-south-1" ]; then
          # ap-southeast-3 & ap-south-1 only use ds-office
          rs_ipset_name="o11y-${env}-regional-${region}-ds-office"
        fi
        update_ipset ${region} ${rs_ipset_name} ${rsIpWhitelist}
      fi
      if [ -n "${dsIpWhitelist}" ]; then
        ds_ipset_name="o11y-${env}-regional-${region}-ds-office"
        if [ "${region}" == "us-west-2" ]; then
          # us-west-2 has a special ds ipset name
          ds_ipset_name="o11y-${env}-regional-${region}-prometheus-office"
        fi
        update_ipset ${region} ${ds_ipset_name} ${dsIpWhitelist}
      fi
      exit 0
    fi
  done
fi

### END - Logic for existed region ###

### BEGIN - Logic for region opened by auto-deploy ###

stack_name="auto-deploy-${cluster}-ingress"
template_file="/tmp/${stack_name}.yaml"
aws cloudformation get-template \
  --region ${region} \
  --stack-name ${stack_name} \
  --query 'TemplateBody' \
  --output text > ${template_file}

if [ -n "${rsIpWhitelist}" ]; then
  IFS=',' read -ra IPS <<< "${rsIpWhitelist}"
  for ip in "${IPS[@]}"; do
    yq eval ".Resources.RsIpSet.Properties.Addresses += [\"${ip}\"]" ${template_file} -i
  done
fi
if [ -n "${dsIpWhitelist}" ]; then
  IFS=',' read -ra IPS <<< "${dsIpWhitelist}"
  for ip in "${IPS[@]}"; do
    yq eval ".Resources.OfficeIpSet.Properties.Addresses += [\"${ip}\"]" ${template_file} -i
  done
fi

parameters=$(aws cloudformation describe-stacks \
  --region ${region} \
  --stack-name ${stack_name} \
  --query 'Stacks[0].Parameters' \
  --output json)

aws cloudformation update-stack \
  --region ${region} \
  --stack-name ${stack_name} \
  --template-body file://${template_file} \
  --capabilities CAPABILITY_NAMED_IAM \
  --parameters "${parameters}"

aws cloudformation wait stack-update-complete --region ${region} --stack-name ${stack_name}

rm -f ${template_file}

### END - Logic for region opened by auto-deploy ###