#!/bin/bash
# Please ensure 'jq' and AWS CLI v2 are installed before executing this script.

set -euo pipefail

output_file="../src/autoconfigure/instance_throughput.rs"
timestamp="$(date --utc +%FT%TZ)"

# Fetch all enabled regions
regions=(us-east-1 us-east-2  us-west-2)

echo "Querying ${#regions[@]} regions: ${regions[*]}"

# Special case handling for Mult-NIC instances
# i.e. dll.24xlarge  --> 4x 100 Gbps
declare -r -A THROUGHPUT_OVERRIDE=(
    ["dl1.24xlarge"]=400
    ["p4d.24xlarge"]=400
    ["p4de.24xlarge"]=400
    ["trn1.32xlarge"]=800
    ["trn1n.32xlarge"]=1600
    ["trn2.48xlarge"]=32000
)

# Create temporary directory for region results
temp_dir=$(mktemp -d)
trap 'rm -rf "$temp_dir"' EXIT

# Function to query a single region
query_region() {
    local region=$1
    local temp_file="${temp_dir}/${region}.json"

    echo "Querying region: ${region}"
    aws ec2 describe-instance-types \
        --filters "Name=instance-type,Values=*" \
        --query "InstanceTypes[].[InstanceType, NetworkInfo.NetworkPerformance]" \
        --region "${region}" \
        --output json > "${temp_file}" 2>/dev/null || {
            echo "Failed to query region ${region}"
            return 1
        }

    echo "Completed query for ${region}"
    echo "${temp_file}"
}

echo "Querying regions..."
for region in "${regions[@]}"; do
    query_region "${region}"
done

# Check if we have any results
if [ ! "$(ls -A ${temp_dir})" ]; then
    echo "Error: No data was retrieved from any region"
    exit 1
fi

# Merge results from all regions
echo "Merging results..."
merged_json=$(jq -s 'add | unique_by(.[0])' "${temp_dir}"/*.json)

# Start generating Rust code
{
    echo "//! DO NOT EDIT THIS FILE DIRECTLY"
    echo "//! This file is auto-generated by \`../../scripts/network_performance.sh\`"
    echo "//!"
    echo "//! Generated from AWS EC2 API"
    echo "//! Regions: ${regions[*]}"
    echo "//! Timestamp: ${timestamp}"
    echo ""
    echo "/// Instance throughput for an EC2 instance type, if known."
    echo "pub fn get_instance_throughput(instance_type: &str) -> Option<f64> {"
    echo "    match instance_type {"
} > "${output_file}"

# Process merged results
echo "${merged_json}" | \
    jq -c 'sort_by(.[0]) | .[]' | \
    while read line; do
        instance_type=$(echo $line | cut -d ',' -f1 | sed 's/\[//' | sed 's/\"//g')

        if [ -v THROUGHPUT_OVERRIDE[$instance_type] ]; then
            throughput=${THROUGHPUT_OVERRIDE[$instance_type]}
            printf "        \"%s\" => Some(%.2f),\n" "${instance_type}" "${throughput}" >> "${output_file}"
        else
            throughput=$(echo $line | cut -d ',' -f2 | sed 's/\s*[a-zA-Z]\s*//g' | sed 's/\]//' | sed 's/\"//g')
            if [[ -z "$throughput" ]]; then
                printf "        \"%s\" => None,\n" "${instance_type}" >> "${output_file}"
            else
                printf "        \"%s\" => Some(%.2f),\n" "${instance_type}" "${throughput}" >> "${output_file}"
            fi
        fi
    done

# Close the Rust code
{
    echo "        _ => None,"
    echo "    }"
    echo "}"
} >> "${output_file}"

echo "Generated ${output_file}"