//! AWS EC2 SDK function wrappers

use super::{METRICS_PORT, SYSTEM_PORT};
use crate::aws::{
    PortConfig,
    utils::{
        DEPLOYER_MAX_PORT, DEPLOYER_MIN_PORT, DEPLOYER_PROTOCOL, MAX_SDK_ATTEMPTS, RETRY_INTERVAL,
        SDK_INITIAL_BACKOFF, SDK_MAX_BACKOFF, exact_cidr,
    },
};
pub use aws_config::Region;
use aws_config::{BehaviorVersion, retry::RetryConfig};
pub use aws_sdk_ec2::{
    Client as Ec2Client,
    types::{InstanceType, IpPermission, IpRange, UserIdGroupPair, VolumeType},
};
use aws_sdk_ec2::{
    Error as Ec2Error,
    error::{BuildError, ProvideErrorMetadata as _, SdkError},
    operation::run_instances::RunInstancesError,
    primitives::Blob,
    types::{
        BlockDeviceMapping, EbsBlockDevice, EphemeralNvmeSupport, Filter, InstanceStateName,
        InstanceTypeInfo, ResourceType, SecurityGroup, SummaryStatus, Tag, TagSpecification,
        VpcPeeringConnectionStateReasonCode,
    },
};
use std::{
    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
    time::Duration,
};
use tokio::time::sleep;
use tracing::{debug, warn};

#[cfg(not(test))]
const LAUNCH_RETRY_INTERVAL: Duration = RETRY_INTERVAL;
#[cfg(test)]
const LAUNCH_RETRY_INTERVAL: Duration = Duration::ZERO;

type LaunchSdkError = SdkError<RunInstancesError>;

// RunInstances disables SDK retries below, so its owner retains the SDK's standard transient
// service categories while handling capacity responses at the availability-zone level.
const FATAL_LAUNCH_ERROR_CODE_PREFIXES: &[&str] = &[
    "UnauthorizedOperation",
    "OptInRequired",
    "VcpuLimitExceeded",
    "InstanceLimitExceeded",
    "MaxSpotInstanceCountExceeded",
    "VolumeLimitExceeded",
    "InvalidParameterValue",
    "InvalidAMIID",
    "InvalidSubnetID",
    "InvalidGroup",
    "InvalidKeyPair",
];
const RETRYABLE_LAUNCH_ERROR_CODES: &[&str] = &[
    "Throttling",
    "ThrottlingException",
    "ThrottledException",
    "RequestThrottledException",
    "TooManyRequestsException",
    "ProvisionedThroughputExceededException",
    "TransactionInProgressException",
    "RequestLimitExceeded",
    "BandwidthLimitExceeded",
    "LimitExceededException",
    "RequestThrottled",
    "SlowDown",
    "PriorRequestNotComplete",
    "EC2ThrottledException",
    "RequestTimeout",
    "RequestTimeoutException",
];
const RETRYABLE_LAUNCH_STATUS_CODES: &[u16] = &[500, 502, 503, 504];

/// Creates an EC2 client for the specified AWS region
pub async fn create_client(region: Region) -> Ec2Client {
    let retry = aws_config::retry::RetryConfig::adaptive()
        .with_max_attempts(MAX_SDK_ATTEMPTS)
        .with_initial_backoff(SDK_INITIAL_BACKOFF)
        .with_max_backoff(SDK_MAX_BACKOFF)
        .with_reconnect_mode(aws_sdk_ec2::config::retry::ReconnectMode::ReconnectOnTransientError);
    let config = aws_config::defaults(BehaviorVersion::v2026_01_12())
        .region(region)
        .retry_config(retry)
        .load()
        .await;
    Ec2Client::new(&config)
}

/// Imports an SSH public key into the specified region
pub async fn import_key_pair(
    client: &Ec2Client,
    key_name: &str,
    public_key: &str,
) -> Result<(), Ec2Error> {
    let blob = Blob::new(public_key.as_bytes());
    client
        .import_key_pair()
        .key_name(key_name)
        .public_key_material(blob)
        .send()
        .await?;
    Ok(())
}

/// Deletes an SSH key pair from the specified region
pub async fn delete_key_pair(client: &Ec2Client, key_name: &str) -> Result<(), Ec2Error> {
    client.delete_key_pair().key_name(key_name).send().await?;
    Ok(())
}

async fn describe_instance_type(
    client: &Ec2Client,
    instance_type: &str,
) -> Result<InstanceTypeInfo, Ec2Error> {
    let response = client
        .describe_instance_types()
        .instance_types(InstanceType::try_parse(instance_type).expect("invalid instance type"))
        .send()
        .await?;

    response
        .instance_types
        .and_then(|types| types.into_iter().next())
        .ok_or_else(|| {
            Ec2Error::from(BuildError::other(format!(
                "instance type {instance_type} not found"
            )))
        })
}

/// Detects the architecture of an instance type using the AWS API
pub(crate) async fn detect_architecture(
    client: &Ec2Client,
    instance_type: &str,
) -> Result<super::Architecture, Ec2Error> {
    let instance_info = describe_instance_type(client, instance_type).await?;

    let architectures = instance_info
        .processor_info
        .and_then(|p| p.supported_architectures)
        .unwrap_or_default();

    // EC2 instance types only support one architecture (e.g., t4g.* = arm64, t3.* = x86_64),
    // so the check order here doesn't matter in practice.
    if architectures.iter().any(|a| a.as_ref() == "arm64") {
        Ok(super::Architecture::Arm64)
    } else if architectures.iter().any(|a| a.as_ref() == "x86_64") {
        Ok(super::Architecture::X86_64)
    } else {
        Err(Ec2Error::from(BuildError::other(format!(
            "instance type {instance_type} has no supported architecture"
        ))))
    }
}

/// Checks whether an instance type exposes EC2 NVMe instance-store devices.
pub(crate) async fn supports_nvme_instance_storage(
    client: &Ec2Client,
    instance_type: &str,
) -> Result<bool, Ec2Error> {
    let instance_info = describe_instance_type(client, instance_type).await?;

    Ok(instance_info.instance_storage_supported().unwrap_or(false)
        && instance_info
            .instance_storage_info()
            .and_then(|storage| storage.nvme_support())
            .is_some_and(|support| {
                matches!(
                    support,
                    EphemeralNvmeSupport::Required | EphemeralNvmeSupport::Supported
                )
            }))
}

/// Finds the latest Ubuntu 24.04 AMI for the given architecture in the region
pub(crate) async fn find_latest_ami(
    client: &Ec2Client,
    architecture: super::Architecture,
) -> Result<String, Ec2Error> {
    let arch = architecture.as_str();
    let resp = client
        .describe_images()
        .filters(
            Filter::builder()
                .name("name")
                .values(format!(
                    "ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-{arch}-server-*"
                ))
                .build(),
        )
        .filters(
            Filter::builder()
                .name("root-device-type")
                .values("ebs")
                .build(),
        )
        .owners("099720109477") // Canonical's AWS account ID
        .send()
        .await?;
    let mut images = resp.images.unwrap_or_default();
    if images.is_empty() {
        return Err(Ec2Error::from(BuildError::other(
            "No matching AMI found".to_string(),
        )));
    }
    images.sort_by(|a, b| b.creation_date().cmp(&a.creation_date()));
    let latest_ami = images[0].image_id().unwrap();
    Ok(latest_ami.to_string())
}

/// Creates a VPC with the specified CIDR block and tag
pub async fn create_vpc(
    client: &Ec2Client,
    cidr_block: &str,
    tag: &str,
) -> Result<String, Ec2Error> {
    let resp = client
        .create_vpc()
        .cidr_block(cidr_block)
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::Vpc)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    Ok(resp.vpc.unwrap().vpc_id.unwrap())
}

/// Creates an Internet Gateway and attaches it to the specified VPC
pub async fn create_and_attach_igw(
    client: &Ec2Client,
    vpc_id: &str,
    tag: &str,
) -> Result<String, Ec2Error> {
    let igw_resp = client
        .create_internet_gateway()
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::InternetGateway)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    let igw_id = igw_resp
        .internet_gateway
        .unwrap()
        .internet_gateway_id
        .unwrap();
    client
        .attach_internet_gateway()
        .internet_gateway_id(&igw_id)
        .vpc_id(vpc_id)
        .send()
        .await?;
    Ok(igw_id)
}

/// Creates a route table for the VPC and sets up a default route to the Internet Gateway
pub async fn create_route_table(
    client: &Ec2Client,
    vpc_id: &str,
    igw_id: &str,
    tag: &str,
) -> Result<String, Ec2Error> {
    let rt_resp = client
        .create_route_table()
        .vpc_id(vpc_id)
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::RouteTable)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    let rt_id = rt_resp.route_table.unwrap().route_table_id.unwrap();
    client
        .create_route()
        .route_table_id(&rt_id)
        .destination_cidr_block("0.0.0.0/0")
        .gateway_id(igw_id)
        .send()
        .await?;
    Ok(rt_id)
}

/// Creates a subnet within the VPC and associates it with the route table
pub async fn create_subnet(
    client: &Ec2Client,
    vpc_id: &str,
    route_table_id: &str,
    subnet_cidr: &str,
    availability_zone: &str,
    tag: &str,
) -> Result<String, Ec2Error> {
    let subnet_resp = client
        .create_subnet()
        .vpc_id(vpc_id)
        .cidr_block(subnet_cidr)
        .availability_zone(availability_zone)
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::Subnet)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    let subnet_id = subnet_resp.subnet.unwrap().subnet_id.unwrap();
    client
        .associate_route_table()
        .route_table_id(route_table_id)
        .subnet_id(&subnet_id)
        .send()
        .await?;
    Ok(subnet_id)
}

/// Creates a security group for the monitoring instance with access from the deployer IP
pub async fn create_security_group_monitoring(
    client: &Ec2Client,
    vpc_id: &str,
    deployer_ip: &str,
    tag: &str,
) -> Result<String, Ec2Error> {
    let sg_resp = client
        .create_security_group()
        .group_name(tag)
        .description("Security group for monitoring instance")
        .vpc_id(vpc_id)
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::SecurityGroup)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    let sg_id = sg_resp.group_id.unwrap();
    client
        .authorize_security_group_ingress()
        .group_id(&sg_id)
        .ip_permissions(
            IpPermission::builder()
                .ip_protocol(DEPLOYER_PROTOCOL)
                .from_port(DEPLOYER_MIN_PORT)
                .to_port(DEPLOYER_MAX_PORT)
                .ip_ranges(IpRange::builder().cidr_ip(exact_cidr(deployer_ip)).build())
                .build(),
        )
        .send()
        .await?;
    Ok(sg_id)
}

/// Creates a security group for binary instances with access from deployer and custom ports
/// Note: monitoring IP rules are added separately via `add_monitoring_ingress` after monitoring instance launches
pub async fn create_security_group_binary(
    client: &Ec2Client,
    vpc_id: &str,
    deployer_ip: &str,
    tag: &str,
    ports: &[PortConfig],
) -> Result<String, Ec2Error> {
    let sg_resp = client
        .create_security_group()
        .group_name(format!("{tag}-binary"))
        .description("Security group for binary instances")
        .vpc_id(vpc_id)
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::SecurityGroup)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    let sg_id = sg_resp.group_id.unwrap();
    let mut builder = client
        .authorize_security_group_ingress()
        .group_id(&sg_id)
        .ip_permissions(
            IpPermission::builder()
                .ip_protocol(DEPLOYER_PROTOCOL)
                .from_port(DEPLOYER_MIN_PORT)
                .to_port(DEPLOYER_MAX_PORT)
                .ip_ranges(IpRange::builder().cidr_ip(exact_cidr(deployer_ip)).build())
                .build(),
        );
    for port in ports {
        builder = builder.ip_permissions(
            IpPermission::builder()
                .ip_protocol(&port.protocol)
                .from_port(port.port as i32)
                .to_port(port.port as i32)
                .ip_ranges(IpRange::builder().cidr_ip(&port.cidr).build())
                .build(),
        );
    }
    builder.send().await?;
    Ok(sg_id)
}

/// Adds monitoring IP ingress rules to a binary security group for Prometheus scraping
pub async fn add_monitoring_ingress(
    client: &Ec2Client,
    sg_id: &str,
    monitoring_ip: &str,
) -> Result<(), Ec2Error> {
    client
        .authorize_security_group_ingress()
        .group_id(sg_id)
        .ip_permissions(
            IpPermission::builder()
                .ip_protocol("tcp")
                .from_port(METRICS_PORT as i32)
                .to_port(METRICS_PORT as i32)
                .ip_ranges(
                    IpRange::builder()
                        .cidr_ip(exact_cidr(monitoring_ip))
                        .build(),
                )
                .build(),
        )
        .ip_permissions(
            IpPermission::builder()
                .ip_protocol("tcp")
                .from_port(SYSTEM_PORT as i32)
                .to_port(SYSTEM_PORT as i32)
                .ip_ranges(
                    IpRange::builder()
                        .cidr_ip(exact_cidr(monitoring_ip))
                        .build(),
                )
                .build(),
        )
        .send()
        .await?;
    Ok(())
}

/// Parses a configured EBS storage class.
pub(crate) fn parse_storage_class(
    target: &str,
    storage_class: &str,
) -> Result<VolumeType, super::Error> {
    VolumeType::try_parse(storage_class).map_err(|_| super::Error::InvalidStorageClass {
        target: target.to_string(),
        storage_class: storage_class.to_string(),
    })
}

/// Validates configured EBS options.
///
/// Source docs for EBS request-side limits:
/// - IOPS and throughput request fields:
///   https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_EbsBlockDevice.html
/// - gp3 size, IOPS, and throughput ratios:
///   https://docs.aws.amazon.com/ebs/latest/userguide/general-purpose.html
/// - io1/io2 size-to-IOPS ratios:
///   https://docs.aws.amazon.com/ebs/latest/userguide/provisioned-iops.html
pub(crate) fn validate_storage_options(
    target: &str,
    storage_class: &VolumeType,
    storage_size: i32,
    storage_iops: Option<i32>,
    storage_throughput: Option<i32>,
) -> Result<(), super::Error> {
    // Provisioned IOPS SSD volumes require an explicit IOPS value at launch.
    if storage_iops.is_none() && matches!(storage_class, VolumeType::Io1 | VolumeType::Io2) {
        return Err(super::Error::MissingStorageIops {
            target: target.to_string(),
            storage_class: storage_class.as_str().to_string(),
        });
    }

    // Request IOPS is limited by both the EbsBlockDevice range and the
    // volume-type-specific storage size ratio.
    if let Some(storage_iops) = storage_iops {
        let storage_size = i64::from(storage_size);
        let storage_iops = i64::from(storage_iops);
        let valid = match storage_class {
            VolumeType::Gp3 => {
                let max_iops = 80_000.min(3_000.max(storage_size * 500));
                (3_000..=max_iops).contains(&storage_iops)
            }
            VolumeType::Io1 => {
                let max_iops = 64_000.min(storage_size * 50);
                (100..=max_iops).contains(&storage_iops)
            }
            VolumeType::Io2 => {
                let max_iops = 256_000.min(storage_size * 1_000);
                (100..=max_iops).contains(&storage_iops)
            }
            _ => false,
        };
        if !valid {
            return Err(super::Error::InvalidStorageIops {
                target: target.to_string(),
                storage_class: storage_class.as_str().to_string(),
                storage_iops: storage_iops as i32,
            });
        }
    }

    // EbsBlockDevice throughput is a gp3-only request field. gp3 throughput
    // is capped at 0.25 MiB/s per effective provisioned IOPS.
    match (storage_throughput, storage_class) {
        (Some(storage_throughput), _) if !(125..=2_000).contains(&storage_throughput) => {
            Err(super::Error::InvalidStorageThroughput {
                target: target.to_string(),
                storage_throughput,
            })
        }
        (Some(_), storage_class) if !matches!(storage_class, VolumeType::Gp3) => {
            Err(super::Error::UnsupportedStorageThroughput {
                target: target.to_string(),
                storage_class: storage_class.as_str().to_string(),
            })
        }
        (Some(storage_throughput), VolumeType::Gp3)
            if storage_throughput > storage_iops.unwrap_or(3_000) / 4 =>
        {
            Err(super::Error::InvalidStorageThroughput {
                target: target.to_string(),
                storage_throughput,
            })
        }
        _ => Ok(()),
    }
}

/// Attempts to launch EC2 instances. May fail on transient errors or rate limits.
#[allow(clippy::too_many_arguments)]
async fn try_launch_instances(
    client: &Ec2Client,
    ami_id: &str,
    instance_type: InstanceType,
    storage_size: i32,
    storage_class: VolumeType,
    storage_iops: Option<i32>,
    storage_throughput: Option<i32>,
    key_name: &str,
    subnet_id: &str,
    sg_id: &str,
    count: i32,
    name: &str,
    tag: &str,
    client_token: &str,
) -> Result<Vec<String>, LaunchSdkError> {
    // Build the root EBS mapping with optional provisioned performance settings.
    let mut ebs = EbsBlockDevice::builder()
        .volume_size(storage_size)
        .volume_type(storage_class)
        .delete_on_termination(true);
    if let Some(storage_iops) = storage_iops {
        ebs = ebs.iops(storage_iops);
    }
    if let Some(storage_throughput) = storage_throughput {
        ebs = ebs.throughput(storage_throughput);
    }

    // Send one request because `launch_instances` owns subnet selection, retry classification,
    // and retry cadence.
    let resp = client
        .run_instances()
        .image_id(ami_id)
        .instance_type(instance_type)
        .key_name(key_name)
        .min_count(count)
        .max_count(count)
        .client_token(client_token)
        .network_interfaces(
            aws_sdk_ec2::types::InstanceNetworkInterfaceSpecification::builder()
                .associate_public_ip_address(true)
                .device_index(0)
                .subnet_id(subnet_id)
                .groups(sg_id)
                .build(),
        )
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::Instance)
                .set_tags(Some(vec![
                    Tag::builder().key("deployer").value(tag).build(),
                    Tag::builder().key("name").value(name).build(),
                ]))
                .build(),
        )
        .block_device_mappings(
            BlockDeviceMapping::builder()
                .device_name("/dev/sda1")
                .ebs(ebs.build())
                .build(),
        )
        .customize()
        .config_override(aws_sdk_ec2::config::Builder::new().retry_config(RetryConfig::disabled()))
        .send()
        .await?;
    Ok(resp
        .instances
        .unwrap()
        .into_iter()
        .map(|i| i.instance_id.unwrap())
        .collect())
}

/// Extracts the structured EC2 code from a RunInstances service error.
fn launch_error_code(error: &LaunchSdkError) -> Option<&str> {
    match error {
        SdkError::ServiceError(context) => context.err().code(),
        _ => None,
    }
}

/// Checks if an EC2 error may resolve after capacity changes in the region.
fn is_capacity_error(error: &LaunchSdkError) -> bool {
    launch_error_code(error) == Some("InsufficientInstanceCapacity")
}

/// Checks if an EC2 error makes a subnet unusable for the remainder of this launch.
fn is_subnet_unavailable_error(error: &LaunchSdkError) -> bool {
    launch_error_code(error) == Some("InsufficientFreeAddressesInSubnet")
}

/// Checks if an EC2 error code belongs to a fatal authorization, configuration, or quota family.
/// Some EC2 code families append a subtype suffix, such as `InvalidAMIID.NotFound`.
fn is_fatal_launch_error_code(code: &str) -> bool {
    FATAL_LAUNCH_ERROR_CODE_PREFIXES
        .iter()
        .any(|prefix| code.starts_with(prefix))
}

/// Checks if retrying the same idempotent RunInstances request may resolve the failure.
fn is_retryable_launch_error(error: &LaunchSdkError) -> bool {
    match error {
        SdkError::TimeoutError(_) | SdkError::ResponseError(_) => true,
        SdkError::DispatchFailure(context) => {
            context.is_io() || context.is_timeout() || context.as_other().is_some()
        }
        SdkError::ServiceError(context) => {
            let code = context.err().code();
            !code.is_some_and(is_fatal_launch_error_code)
                && (code.is_some_and(|code| RETRYABLE_LAUNCH_ERROR_CODES.contains(&code))
                    || RETRYABLE_LAUNCH_STATUS_CODES.contains(&context.raw().status().as_u16()))
        }
        _ => false,
    }
}

/// Launches EC2 instances with specified configurations.
/// Filters subnets to those supporting the instance type, distributes across them starting at
/// `start_idx`, and retries complete eligible-AZ scans on capacity errors.
#[allow(clippy::too_many_arguments)]
pub async fn launch_instances(
    client: &Ec2Client,
    ami_id: &str,
    instance_type: InstanceType,
    storage_size: i32,
    storage_class: VolumeType,
    storage_iops: Option<i32>,
    storage_throughput: Option<i32>,
    key_name: &str,
    subnets: &[(String, String)],
    az_support: &BTreeMap<String, BTreeSet<String>>,
    start_idx: usize,
    sg_id: &str,
    count: i32,
    name: &str,
    tag: &str,
) -> Result<(Vec<String>, String), super::Error> {
    validate_storage_options(
        name,
        &storage_class,
        storage_size,
        storage_iops,
        storage_throughput,
    )?;

    // Filter to subnets in AZs that support this instance type
    let instance_type_str = instance_type.to_string();
    let eligible: Vec<(&str, &str)> = subnets
        .iter()
        .filter(|(az, _)| {
            az_support
                .get(az)
                .is_some_and(|types| types.contains(&instance_type_str))
        })
        .map(|(az, subnet_id)| (az.as_str(), subnet_id.as_str()))
        .collect();
    if eligible.is_empty() {
        return Err(super::Error::UnsupportedInstanceType(instance_type_str));
    }

    let len = eligible.len();
    let mut last_error = None;
    let mut unavailable = vec![false; len];
    let mut scan = 0u64;
    loop {
        // Each scan probes every usable subnet once. Capacity errors keep a subnet eligible for the
        // next scan, while address exhaustion permanently removes it.
        scan = scan.saturating_add(1);
        let mut retry_capacity = false;
        for i in 0..len {
            let eligible_index = (start_idx + i) % len;
            if unavailable[eligible_index] {
                continue;
            }
            let (az, subnet_id) = eligible[eligible_index];

            // A probe owns one client token across ambiguous retries. Moving to another subnet
            // changes the request parameters and starts a new probe with a new token.
            let client_token = uuid::Uuid::new_v4().to_string();
            loop {
                match try_launch_instances(
                    client,
                    ami_id,
                    instance_type.clone(),
                    storage_size,
                    storage_class.clone(),
                    storage_iops,
                    storage_throughput,
                    key_name,
                    subnet_id,
                    sg_id,
                    count,
                    name,
                    tag,
                    &client_token,
                )
                .await
                {
                    Ok(ids) => return Ok((ids, az.to_string())),
                    Err(e) if is_capacity_error(&e) => {
                        retry_capacity = true;
                        warn!(
                            name = name,
                            az,
                            scan,
                            error = %e,
                            "insufficient instance capacity, trying next subnet"
                        );
                        last_error = Some(e.into());
                        break;
                    }
                    Err(e) if is_subnet_unavailable_error(&e) => {
                        unavailable[eligible_index] = true;
                        debug!(
                            name = name,
                            az,
                            error = %e,
                            "subnet unavailable, trying next subnet"
                        );
                        last_error = Some(e.into());
                        break;
                    }
                    Err(e) if !is_retryable_launch_error(&e) => {
                        return Err(super::Error::AwsEc2(e.into()));
                    }
                    Err(e) => {
                        debug!(
                            name = name,
                            error = %e,
                            "launch_instances failed, retrying"
                        );
                        sleep(LAUNCH_RETRY_INTERVAL).await;
                    }
                }
            }
        }

        // Start another scan only after a capacity failure. All other completed scans are terminal.
        if !retry_capacity {
            break;
        }
        debug!(
            name,
            scan, "capacity unavailable in every usable AZ, waiting before retry"
        );
        sleep(LAUNCH_RETRY_INTERVAL).await;
    }

    Err(last_error.map_or(super::Error::NoSubnetsAvailable, super::Error::AwsEc2))
}

/// Waits for instances to reach the "running" state and returns their public IPs
/// in the same order as the input instance IDs.
pub async fn wait_for_instances_running(
    client: &Ec2Client,
    instance_ids: &[String],
) -> Result<Vec<String>, Ec2Error> {
    // Track discovered IPs to avoid re-polling running instances
    let mut discovered_ips: HashMap<String, String> = HashMap::new();
    let mut pending_ids: HashSet<String> = instance_ids.iter().cloned().collect();
    let mut attempt = 0u32;
    loop {
        // Only query instances that haven't been discovered yet
        let query_ids: Vec<String> = pending_ids.iter().cloned().collect();
        let resp = match client
            .describe_instances()
            .set_instance_ids(Some(query_ids))
            .send()
            .await
        {
            Ok(resp) => {
                attempt = 0;
                resp
            }
            Err(e) => {
                attempt = attempt.saturating_add(1);
                debug!(
                    pending = pending_ids.len(),
                    attempt = attempt,
                    error = %e,
                    "describe_instances failed, retrying"
                );
                sleep(RETRY_INTERVAL).await;
                continue;
            }
        };

        // Check each instance and record those that are running with IPs
        for reservation in resp.reservations.unwrap_or_default() {
            for instance in reservation.instances.unwrap_or_default() {
                let id = match instance.instance_id {
                    Some(id) => id,
                    None => continue,
                };
                let is_running = instance.state.as_ref().and_then(|s| s.name.as_ref())
                    == Some(&InstanceStateName::Running);
                if is_running {
                    if let Some(ip) = instance.public_ip_address {
                        discovered_ips.insert(id.clone(), ip);
                        pending_ids.remove(&id);
                    }
                }
            }
        }

        // Return once all instances are discovered
        if pending_ids.is_empty() {
            return Ok(instance_ids
                .iter()
                .map(|id| discovered_ips.remove(id).unwrap())
                .collect());
        }

        // Try again after a delay
        sleep(RETRY_INTERVAL).await;
    }
}

pub async fn wait_for_instances_ready(
    client: &Ec2Client,
    instance_ids: &[String],
) -> Result<(), Ec2Error> {
    loop {
        // Ask for instance status
        let Ok(resp) = client
            .describe_instance_status()
            .set_instance_ids(Some(instance_ids.to_vec()))
            .include_all_instances(true) // Include instances regardless of state
            .send()
            .await
        else {
            sleep(RETRY_INTERVAL).await;
            continue;
        };

        // Confirm all are ready
        let statuses = resp.instance_statuses.unwrap_or_default();
        let all_ready = statuses.iter().all(|s| {
            s.instance_state.as_ref().unwrap().name.as_ref().unwrap() == &InstanceStateName::Running
                && s.system_status.as_ref().unwrap().status.as_ref().unwrap() == &SummaryStatus::Ok
                && s.instance_status.as_ref().unwrap().status.as_ref().unwrap()
                    == &SummaryStatus::Ok
        });
        if !all_ready {
            sleep(RETRY_INTERVAL).await;
            continue;
        }
        return Ok(());
    }
}

/// Retrieves the private IP address of an instance
pub async fn get_private_ip(client: &Ec2Client, instance_id: &str) -> Result<String, Ec2Error> {
    let resp = client
        .describe_instances()
        .instance_ids(instance_id)
        .send()
        .await?;
    let reservations = resp.reservations.unwrap();
    let instance = &reservations[0].instances.as_ref().unwrap()[0];
    Ok(instance.private_ip_address.as_ref().unwrap().clone())
}

/// Creates a VPC peering connection between two VPCs
pub async fn create_vpc_peering_connection(
    client: &Ec2Client,
    requester_vpc_id: &str,
    peer_vpc_id: &str,
    peer_region: &str,
    tag: &str,
) -> Result<String, Ec2Error> {
    let resp = client
        .create_vpc_peering_connection()
        .vpc_id(requester_vpc_id)
        .peer_vpc_id(peer_vpc_id)
        .peer_region(peer_region)
        .tag_specifications(
            TagSpecification::builder()
                .resource_type(ResourceType::VpcPeeringConnection)
                .tags(Tag::builder().key("deployer").value(tag).build())
                .build(),
        )
        .send()
        .await?;
    Ok(resp
        .vpc_peering_connection
        .unwrap()
        .vpc_peering_connection_id
        .unwrap())
}

/// Waits for a VPC peering connection to reach the "pending-acceptance" state
pub async fn wait_for_vpc_peering_connection(
    client: &Ec2Client,
    peer_id: &str,
) -> Result<(), Ec2Error> {
    loop {
        if let Ok(resp) = client
            .describe_vpc_peering_connections()
            .vpc_peering_connection_ids(peer_id)
            .send()
            .await
        {
            if let Some(connections) = resp.vpc_peering_connections {
                if let Some(connection) = connections.first() {
                    if connection.status.as_ref().unwrap().code
                        == Some(VpcPeeringConnectionStateReasonCode::PendingAcceptance)
                    {
                        return Ok(());
                    }
                }
            }
        }
        sleep(Duration::from_secs(2)).await;
    }
}

/// Accepts a VPC peering connection in the peer region
pub async fn accept_vpc_peering_connection(
    client: &Ec2Client,
    peer_id: &str,
) -> Result<(), Ec2Error> {
    client
        .accept_vpc_peering_connection()
        .vpc_peering_connection_id(peer_id)
        .send()
        .await?;
    Ok(())
}

/// Adds a route to a route table for VPC peering
pub async fn add_route(
    client: &Ec2Client,
    route_table_id: &str,
    destination_cidr: &str,
    peer_id: &str,
) -> Result<(), Ec2Error> {
    client
        .create_route()
        .route_table_id(route_table_id)
        .destination_cidr_block(destination_cidr)
        .vpc_peering_connection_id(peer_id)
        .send()
        .await?;
    Ok(())
}

/// Finds VPC peering connections by deployer tag
pub async fn find_vpc_peering_by_tag(
    client: &Ec2Client,
    tag: &str,
) -> Result<Vec<String>, Ec2Error> {
    let resp = client
        .describe_vpc_peering_connections()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .vpc_peering_connections
        .unwrap_or_default()
        .into_iter()
        .map(|p| p.vpc_peering_connection_id.unwrap())
        .collect())
}

/// Deletes a VPC peering connection
pub async fn delete_vpc_peering(client: &Ec2Client, peering_id: &str) -> Result<(), Ec2Error> {
    client
        .delete_vpc_peering_connection()
        .vpc_peering_connection_id(peering_id)
        .send()
        .await?;
    Ok(())
}

/// Waits for a VPC peering connection to be deleted
pub async fn wait_for_vpc_peering_deletion(
    ec2_client: &Ec2Client,
    peer_id: &str,
) -> Result<(), Ec2Error> {
    loop {
        let resp = ec2_client
            .describe_vpc_peering_connections()
            .vpc_peering_connection_ids(peer_id)
            .send()
            .await?;
        if let Some(connections) = resp.vpc_peering_connections {
            if let Some(connection) = connections.first() {
                if connection.status.as_ref().unwrap().code
                    == Some(VpcPeeringConnectionStateReasonCode::Deleted)
                {
                    return Ok(());
                }
            } else {
                return Ok(());
            }
        } else {
            return Ok(());
        }
        sleep(RETRY_INTERVAL).await;
    }
}

/// Finds instances by deployer tag
pub async fn find_instances_by_tag(
    ec2_client: &Ec2Client,
    tag: &str,
) -> Result<Vec<String>, Ec2Error> {
    let resp = ec2_client
        .describe_instances()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .reservations
        .unwrap_or_default()
        .into_iter()
        .flat_map(|r| r.instances.unwrap_or_default())
        .map(|i| i.instance_id.unwrap())
        .collect())
}

/// Terminates specified instances
pub async fn terminate_instances(
    ec2_client: &Ec2Client,
    instance_ids: &[String],
) -> Result<(), Ec2Error> {
    if instance_ids.is_empty() {
        return Ok(());
    }
    ec2_client
        .terminate_instances()
        .set_instance_ids(Some(instance_ids.to_vec()))
        .send()
        .await?;
    Ok(())
}

/// Waits for instances to be terminated
pub async fn wait_for_instances_terminated(
    ec2_client: &Ec2Client,
    instance_ids: &[String],
) -> Result<(), Ec2Error> {
    loop {
        let resp = ec2_client
            .describe_instances()
            .set_instance_ids(Some(instance_ids.to_vec()))
            .send()
            .await?;
        let instances = resp
            .reservations
            .unwrap_or_default()
            .into_iter()
            .flat_map(|r| r.instances.unwrap_or_default())
            .collect::<Vec<_>>();
        if instances.iter().all(|i| {
            i.state.as_ref().unwrap().name.as_ref().unwrap() == &InstanceStateName::Terminated
        }) {
            return Ok(());
        }
        sleep(RETRY_INTERVAL).await;
    }
}

/// Finds security groups by deployer tag
pub async fn find_security_groups_by_tag(
    ec2_client: &Ec2Client,
    tag: &str,
) -> Result<Vec<SecurityGroup>, Ec2Error> {
    let resp = ec2_client
        .describe_security_groups()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .security_groups
        .unwrap_or_default()
        .into_iter()
        .collect())
}

/// Deletes a security group
pub async fn delete_security_group(ec2_client: &Ec2Client, sg_id: &str) -> Result<(), Ec2Error> {
    ec2_client
        .delete_security_group()
        .group_id(sg_id)
        .send()
        .await?;
    Ok(())
}

/// Finds route tables by deployer tag
pub async fn find_route_tables_by_tag(
    ec2_client: &Ec2Client,
    tag: &str,
) -> Result<Vec<String>, Ec2Error> {
    let resp = ec2_client
        .describe_route_tables()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .route_tables
        .unwrap_or_default()
        .into_iter()
        .map(|rt| rt.route_table_id.unwrap())
        .collect())
}

/// Deletes a route table
pub async fn delete_route_table(ec2_client: &Ec2Client, rt_id: &str) -> Result<(), Ec2Error> {
    ec2_client
        .delete_route_table()
        .route_table_id(rt_id)
        .send()
        .await?;
    Ok(())
}

/// Finds Internet Gateways by deployer tag
pub async fn find_igws_by_tag(ec2_client: &Ec2Client, tag: &str) -> Result<Vec<String>, Ec2Error> {
    let resp = ec2_client
        .describe_internet_gateways()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .internet_gateways
        .unwrap_or_default()
        .into_iter()
        .map(|igw| igw.internet_gateway_id.unwrap())
        .collect())
}

/// Finds the VPC ID attached to an Internet Gateway, if any
pub async fn find_vpc_by_igw(
    ec2_client: &Ec2Client,
    igw_id: &str,
) -> Result<Option<String>, Ec2Error> {
    let resp = ec2_client
        .describe_internet_gateways()
        .internet_gateway_ids(igw_id)
        .send()
        .await?;
    Ok(resp
        .internet_gateways
        .and_then(|gws| gws.into_iter().next())
        .and_then(|gw| gw.attachments)
        .and_then(|attachments| attachments.into_iter().next())
        .and_then(|attachment| attachment.vpc_id))
}

/// Returns the set of regions that are enabled for the AWS account
pub async fn get_enabled_regions(ec2_client: &Ec2Client) -> Result<HashSet<String>, Ec2Error> {
    let resp = ec2_client
        .describe_regions()
        .all_regions(true)
        .filters(
            Filter::builder()
                .name("opt-in-status")
                .values("opt-in-not-required")
                .values("opted-in")
                .build(),
        )
        .send()
        .await?;
    Ok(resp
        .regions
        .unwrap_or_default()
        .into_iter()
        .filter_map(|r| r.region_name)
        .collect())
}

/// Detaches an Internet Gateway from a VPC
pub async fn detach_igw(
    ec2_client: &Ec2Client,
    igw_id: &str,
    vpc_id: &str,
) -> Result<(), Ec2Error> {
    ec2_client
        .detach_internet_gateway()
        .internet_gateway_id(igw_id)
        .vpc_id(vpc_id)
        .send()
        .await?;
    Ok(())
}

/// Deletes an Internet Gateway
pub async fn delete_igw(ec2_client: &Ec2Client, igw_id: &str) -> Result<(), Ec2Error> {
    ec2_client
        .delete_internet_gateway()
        .internet_gateway_id(igw_id)
        .send()
        .await?;
    Ok(())
}

/// Finds subnets by deployer tag
pub async fn find_subnets_by_tag(
    ec2_client: &Ec2Client,
    tag: &str,
) -> Result<Vec<String>, Ec2Error> {
    let resp = ec2_client
        .describe_subnets()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .subnets
        .unwrap_or_default()
        .into_iter()
        .map(|subnet| subnet.subnet_id.unwrap())
        .collect())
}

/// Deletes a subnet
pub async fn delete_subnet(ec2_client: &Ec2Client, subnet_id: &str) -> Result<(), Ec2Error> {
    ec2_client
        .delete_subnet()
        .subnet_id(subnet_id)
        .send()
        .await?;
    Ok(())
}

/// Finds VPCs by deployer tag
pub async fn find_vpcs_by_tag(ec2_client: &Ec2Client, tag: &str) -> Result<Vec<String>, Ec2Error> {
    let resp = ec2_client
        .describe_vpcs()
        .filters(Filter::builder().name("tag:deployer").values(tag).build())
        .send()
        .await?;
    Ok(resp
        .vpcs
        .unwrap_or_default()
        .into_iter()
        .map(|vpc| vpc.vpc_id.unwrap())
        .collect())
}

/// Deletes a VPC
pub async fn delete_vpc(ec2_client: &Ec2Client, vpc_id: &str) -> Result<(), Ec2Error> {
    ec2_client.delete_vpc().vpc_id(vpc_id).send().await?;
    Ok(())
}

/// Returns a map of AZ -> set of supported instance types for the given instance types.
pub async fn find_az_instance_support(
    client: &Ec2Client,
    instance_types: &[String],
) -> Result<BTreeMap<String, BTreeSet<String>>, Ec2Error> {
    let offerings = client
        .describe_instance_type_offerings()
        .location_type("availability-zone".into())
        .filters(
            Filter::builder()
                .name("instance-type")
                .set_values(Some(instance_types.to_vec()))
                .build(),
        )
        .send()
        .await?
        .instance_type_offerings
        .unwrap_or_default();

    // Build map of AZ -> supported instance types
    let mut az_to_instance_types: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
    for offering in offerings {
        if let (Some(location), Some(instance_type)) = (
            offering.location,
            offering.instance_type.map(|it| it.to_string()),
        ) {
            az_to_instance_types
                .entry(location)
                .or_default()
                .insert(instance_type);
        }
    }
    if az_to_instance_types.is_empty() {
        return Err(Ec2Error::from(BuildError::other(format!(
            "no availability zone supports any of: {instance_types:?}"
        ))));
    }

    Ok(az_to_instance_types)
}

/// Waits until all network interfaces associated with a security group are deleted
pub async fn wait_for_enis_deleted(ec2_client: &Ec2Client, sg_id: &str) -> Result<(), Ec2Error> {
    loop {
        let resp = ec2_client
            .describe_network_interfaces()
            .filters(Filter::builder().name("group-id").values(sg_id).build())
            .send()
            .await?;
        if resp.network_interfaces.unwrap_or_default().is_empty() {
            return Ok(());
        }
        sleep(RETRY_INTERVAL).await;
    }
}

#[cfg(test)]
mod tests {
    use super::{
        InstanceType, LaunchSdkError, Region, VolumeType, is_retryable_launch_error,
        launch_instances,
    };
    use crate::aws::Error;
    use aws_config::{BehaviorVersion, retry::RetryConfig};
    use aws_sdk_ec2::{
        Client,
        config::{AsyncSleep, Credentials, Sleep},
        error::BuildError,
    };
    use aws_smithy_runtime_api::{
        client::{
            http::{HttpConnector, HttpConnectorFuture, SharedHttpConnector, http_client_fn},
            orchestrator::{HttpRequest, HttpResponse},
            result::ConnectorError,
            retries::ErrorKind,
        },
        http::StatusCode,
    };
    use std::{
        collections::{BTreeMap, BTreeSet},
        sync::{
            Arc, OnceLock,
            atomic::{AtomicUsize, Ordering},
        },
        time::Duration,
    };

    const CAPACITY_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>InsufficientInstanceCapacity</Code><Message>capacity unavailable</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const SUBNET_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>InsufficientFreeAddressesInSubnet</Code><Message>subnet unavailable</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const TRANSIENT_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>InternalError</Code><Message>transient failure</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const THROTTLED_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>RequestLimitExceeded</Code><Message>request limit exceeded</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const UNAUTHORIZED_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>UnauthorizedOperation</Code><Message>not authorized</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const OPT_IN_REQUIRED_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>OptInRequired</Code><Message>region is not enabled</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const VCPU_LIMIT_ERROR: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<Response><Errors><Error><Code>VcpuLimitExceeded</Code><Message>quota exceeded</Message></Error></Errors><RequestID>request-id</RequestID></Response>"#;
    const LAUNCH_SUCCESS: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
<RunInstancesResponse xmlns="http://ec2.amazonaws.com/doc/2016-11-15/"><instancesSet><item><instanceId>i-test</instanceId></item></instancesSet></RunInstancesResponse>"#;

    #[derive(Clone, Copy, Debug)]
    enum ResponseSpec {
        Http { status: u16, body: &'static str },
        IoError,
        TransientOther,
    }

    fn replay_response(status: u16, body: &'static str) -> ResponseSpec {
        ResponseSpec::Http { status, body }
    }

    fn client_token(request_body: &str) -> &str {
        request_body
            .split('&')
            .find_map(|field| field.strip_prefix("ClientToken="))
            .expect("RunInstances should carry a client token")
    }

    fn http_response(status: u16, body: &'static str) -> HttpResponse {
        HttpResponse::new(StatusCode::try_from(status).unwrap(), body.into())
    }

    #[derive(Clone, Debug)]
    struct ReplayConnector {
        responses: Arc<Vec<ResponseSpec>>,
        requests: Arc<AtomicUsize>,
        request_bodies: Arc<Vec<OnceLock<String>>>,
    }

    #[derive(Debug)]
    struct InstantSleep;

    impl AsyncSleep for InstantSleep {
        fn sleep(&self, _duration: std::time::Duration) -> Sleep {
            Sleep::new(std::future::ready(()))
        }
    }

    impl ReplayConnector {
        fn new(responses: Vec<ResponseSpec>) -> Self {
            let request_bodies = (0..responses.len()).map(|_| OnceLock::new()).collect();
            Self {
                responses: Arc::new(responses),
                requests: Arc::new(AtomicUsize::new(0)),
                request_bodies: Arc::new(request_bodies),
            }
        }

        fn request_count(&self) -> usize {
            self.requests.load(Ordering::Relaxed)
        }

        fn request_bodies(&self) -> Vec<String> {
            self.request_bodies
                .iter()
                .filter_map(OnceLock::get)
                .cloned()
                .collect()
        }
    }

    impl HttpConnector for ReplayConnector {
        fn call(&self, request: HttpRequest) -> HttpConnectorFuture {
            let index = self.requests.fetch_add(1, Ordering::Relaxed);
            if let Some(request_body) = self.request_bodies.get(index) {
                request_body
                    .set(
                        String::from_utf8_lossy(request.body().bytes().unwrap_or_default())
                            .into_owned(),
                    )
                    .expect("each scripted request has a unique index");
            }
            let response = self.responses.get(index).copied();
            HttpConnectorFuture::new(async move {
                match response {
                    Some(ResponseSpec::Http { status, body }) => Ok(http_response(status, body)),
                    Some(ResponseSpec::IoError) => Err(ConnectorError::io(
                        std::io::Error::other("scripted connection failure").into(),
                    )),
                    Some(ResponseSpec::TransientOther) => Err(ConnectorError::other(
                        "scripted incomplete response".into(),
                        Some(ErrorKind::TransientError),
                    )),
                    None => Err(ConnectorError::other(
                        "no scripted EC2 response remains".into(),
                        None,
                    )),
                }
            })
        }
    }

    fn client_with_retry(
        responses: Vec<ResponseSpec>,
        retry_config: RetryConfig,
    ) -> (Client, ReplayConnector) {
        let connector = ReplayConnector::new(responses);
        let http_client = http_client_fn({
            let connector = connector.clone();
            move |_, _| SharedHttpConnector::new(connector.clone())
        });
        let config = aws_sdk_ec2::Config::builder()
            .behavior_version(BehaviorVersion::v2026_01_12())
            .region(Region::new("us-east-1"))
            .credentials_provider(Credentials::new(
                "access-key",
                "secret-key",
                None,
                None,
                "test",
            ))
            .retry_config(retry_config)
            .sleep_impl(InstantSleep)
            .http_client(http_client)
            .build();
        (Client::from_conf(config), connector)
    }

    fn client(responses: Vec<ResponseSpec>) -> (Client, ReplayConnector) {
        client_with_retry(responses, RetryConfig::disabled())
    }

    async fn launch(client: &Client) -> Result<(Vec<String>, String), Error> {
        let subnets = vec![
            ("us-east-1a".to_string(), "subnet-a".to_string()),
            ("us-east-1b".to_string(), "subnet-b".to_string()),
        ];
        let instance_type = "c8a.8xlarge";
        let az_support = BTreeMap::from([
            (
                "us-east-1a".to_string(),
                BTreeSet::from([instance_type.to_string()]),
            ),
            (
                "us-east-1b".to_string(),
                BTreeSet::from([instance_type.to_string()]),
            ),
        ]);

        launch_instances(
            client,
            "ami-test",
            InstanceType::from(instance_type),
            10,
            VolumeType::Gp3,
            None,
            None,
            "key-test",
            &subnets,
            &az_support,
            0,
            "sg-test",
            1,
            "instance-test",
            "tag-test",
        )
        .await
    }

    #[test]
    fn request_construction_failure_is_not_retryable() {
        let error = LaunchSdkError::construction_failure(BuildError::other("invalid request"));
        assert!(!is_retryable_launch_error(&error));

        let error = LaunchSdkError::dispatch_failure(ConnectorError::other(
            "credential resolution failed".into(),
            None,
        ));
        assert!(!is_retryable_launch_error(&error));

        let error = LaunchSdkError::timeout_error(std::io::Error::other("timeout"));
        assert!(is_retryable_launch_error(&error));

        let error = LaunchSdkError::dispatch_failure(ConnectorError::io(
            std::io::Error::other("connection reset").into(),
        ));
        assert!(is_retryable_launch_error(&error));
    }

    #[tokio::test]
    async fn capacity_retry_revisits_eligible_subnets() {
        let (client, connector) = client(vec![
            replay_response(400, CAPACITY_ERROR),
            replay_response(400, CAPACITY_ERROR),
            replay_response(200, LAUNCH_SUCCESS),
        ]);

        let (instances, az) = launch(&client)
            .await
            .expect("capacity should be retried after every eligible AZ fails");

        assert_eq!(instances, ["i-test"]);
        assert_eq!(az, "us-east-1a");
        assert_eq!(connector.request_count(), 3);
    }

    #[tokio::test]
    async fn capacity_scan_owns_run_instances_retries() {
        let (client, connector) = client_with_retry(
            vec![
                replay_response(500, CAPACITY_ERROR),
                replay_response(200, LAUNCH_SUCCESS),
            ],
            RetryConfig::standard().with_max_attempts(2),
        );

        let (instances, az) = launch(&client)
            .await
            .expect("capacity failure should advance to the next AZ");

        assert_eq!(instances, ["i-test"]);
        assert_eq!(az, "us-east-1b");
        assert_eq!(connector.request_count(), 2);
        let bodies = connector.request_bodies();
        assert_ne!(client_token(&bodies[0]), client_token(&bodies[1]));
    }

    #[tokio::test]
    async fn same_subnet_retries_reuse_client_token() {
        let (client, connector) = client(vec![
            ResponseSpec::IoError,
            replay_response(200, LAUNCH_SUCCESS),
        ]);

        launch(&client)
            .await
            .expect("a transient failure should retry the same subnet");

        let bodies = connector.request_bodies();
        let client_tokens: Vec<_> = bodies.iter().map(|body| client_token(body)).collect();
        assert_eq!(client_tokens.len(), 2);
        assert_eq!(client_tokens[0], client_tokens[1]);
    }

    #[tokio::test]
    async fn typed_connector_other_retries_same_request() {
        let (client, connector) = client(vec![
            ResponseSpec::TransientOther,
            replay_response(200, LAUNCH_SUCCESS),
        ]);

        launch(&client)
            .await
            .expect("a typed transient connector failure should retry the same subnet");

        let bodies = connector.request_bodies();
        assert_eq!(bodies.len(), 2);
        assert_eq!(client_token(&bodies[0]), client_token(&bodies[1]));
    }

    #[tokio::test]
    async fn permanent_service_errors_are_not_retried() {
        for (status, body) in [(400, UNAUTHORIZED_ERROR), (400, OPT_IN_REQUIRED_ERROR)] {
            let (client, connector) = client(vec![replay_response(status, body)]);
            let result = tokio::time::timeout(Duration::from_millis(50), launch(&client))
                .await
                .expect("permanent service error must return immediately");
            assert!(result.is_err());
            assert_eq!(connector.request_count(), 1);
        }
    }

    #[tokio::test]
    async fn fatal_service_code_overrides_retryable_status() {
        for body in [UNAUTHORIZED_ERROR, OPT_IN_REQUIRED_ERROR, VCPU_LIMIT_ERROR] {
            let (client, connector) = client(vec![replay_response(500, body)]);

            let result = tokio::time::timeout(Duration::from_millis(50), launch(&client))
                .await
                .expect("a fatal service code must return immediately despite its status");

            assert!(result.is_err());
            assert_eq!(connector.request_count(), 1);
        }
    }

    #[tokio::test]
    async fn transient_service_errors_retry_the_same_request() {
        for (status, body) in [(500, TRANSIENT_ERROR), (400, THROTTLED_ERROR)] {
            let (client, connector) = client(vec![
                replay_response(status, body),
                replay_response(200, LAUNCH_SUCCESS),
            ]);

            launch(&client)
                .await
                .expect("transient service errors should retry the same subnet");

            let bodies = connector.request_bodies();
            assert_eq!(bodies.len(), 2);
            assert_eq!(client_token(&bodies[0]), client_token(&bodies[1]));
        }
    }

    #[tokio::test]
    async fn full_subnet_is_not_retried() {
        let (client, connector) = client(vec![
            replay_response(400, SUBNET_ERROR),
            replay_response(400, CAPACITY_ERROR),
            replay_response(200, LAUNCH_SUCCESS),
        ]);

        let (instances, az) = launch(&client)
            .await
            .expect("the remaining AZ should be retried");

        assert_eq!(instances, ["i-test"]);
        assert_eq!(az, "us-east-1b");
        assert_eq!(connector.request_count(), 3);
    }

    #[tokio::test]
    async fn capacity_retries_until_capacity_returns() {
        let failures = 20;
        let mut events: Vec<_> = (0..failures)
            .map(|_| replay_response(400, CAPACITY_ERROR))
            .collect();
        events.push(replay_response(200, LAUNCH_SUCCESS));
        let (client, connector) = client(events);

        let (instances, az) = launch(&client)
            .await
            .expect("regional capacity exhaustion should remain retryable");

        assert_eq!(instances, ["i-test"]);
        assert_eq!(az, "us-east-1a");
        assert_eq!(connector.request_count(), failures + 1);
    }
}