# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2025 ScyllaDB

"""Spark migrator job submission and monitoring for EMR clusters."""

import gzip
import logging
import os
import tempfile
import time
from dataclasses import dataclass, field

import boto3
import requests
import yaml
from botocore.exceptions import ClientError

LOGGER = logging.getLogger(__name__)

MIGRATOR_GITHUB_REPO = "scylladb/scylla-migrator"
MIGRATOR_JAR_NAME = "scylla-migrator-assembly.jar"

# TODO: Spark 4.x workaround
# The scylla-migrator requires Scala 2.13 which ships with Spark 4.x, but Amazon EMR provides only Spark 3.x
# at the moment (which ships Scala 2.12).
# The workaround installs Spark 4.x via a bootstrap action and uses a wrapper shell script
# (through script-runner.jar) instead of a native Spark step.
# Once Amazon EMR is bundled with Spark 4.x out of the box, this workaround can be removed.
#
# The Spark 4.x tarball is mirrored into the SCT-owned S3 bucket on first use, and the bootstrap pulls it
# from there via `aws s3 cp`. This avoids depending on the Apache live mirror (downloads.apache.org), which
# serves only current releases and returns 404 once a version is superseded. The one-time mirror upload reads
# from the Apache archive (archive.apache.org), which retains every release permanently.
SPARK4_VERSION = "4.0.2"
SPARK4_INSTALL_DIR = "/opt/spark4"
_BOOTSTRAP_SCRIPT_S3_KEY = "scripts/bootstrap-spark4.sh"
_SPARK4_PACKAGE = f"spark-{SPARK4_VERSION}-bin-hadoop3"
_SPARK4_ARCHIVE_URL = f"https://archive.apache.org/dist/spark/spark-{SPARK4_VERSION}/{_SPARK4_PACKAGE}.tgz"
_SPARK4_S3_KEY = f"spark/{_SPARK4_PACKAGE}.tgz"
_RUNNER_SCRIPT_S3_KEY = "scripts/run-migrator.sh"
_VALIDATOR_SCRIPT_S3_KEY = "scripts/run-validator.sh"
_MIGRATOR_MAIN_CLASS = "com.scylladb.migrator.Migrator"
_VALIDATOR_MAIN_CLASS = "com.scylladb.migrator.Validator"
_COMMAND_RUNNER_JAR = "command-runner.jar"
_SPARK_SCYLLA_CONFIG_KEY = "spark.scylla.config"


def _build_bootstrap_script(s3_bucket):
    """Return the Spark 4.x install bootstrap script that pulls the tarball from the SCT S3 mirror."""
    return f"""#!/bin/bash
set -ex
cd /tmp
aws s3 cp "s3://{s3_bucket}/{_SPARK4_S3_KEY}" "{_SPARK4_PACKAGE}.tgz"
tar xzf "{_SPARK4_PACKAGE}.tgz"
sudo mv /tmp/{_SPARK4_PACKAGE} {SPARK4_INSTALL_DIR}
sudo chown -R hadoop:hadoop {SPARK4_INSTALL_DIR}
rm -f /tmp/{_SPARK4_PACKAGE}.tgz
echo "Bootstrap complete: Spark {SPARK4_VERSION} installed at {SPARK4_INSTALL_DIR}"
"""


def _ensure_s3_bucket(s3_client, bucket_name, region_name):
    """Create S3 bucket if it does not already exist.

    Args:
        s3_client: boto3 S3 client.
        bucket_name: Name of the bucket to ensure.
        region_name: AWS region for the bucket.
    """
    try:
        s3_client.head_bucket(Bucket=bucket_name)
        return  # bucket already exists
    except ClientError as exc:
        if exc.response["Error"]["Code"] not in ("404", "NoSuchBucket"):
            raise

    LOGGER.info("Creating S3 bucket %s in %s", bucket_name, region_name)
    create_params = {"Bucket": bucket_name}
    if region_name != "us-east-1":
        create_params["CreateBucketConfiguration"] = {"LocationConstraint": region_name}
    s3_client.create_bucket(**create_params)


# TODO: Spark 4.x workaround
def upload_spark4_tarball_to_s3(s3_bucket, region_name):
    """Mirror the Spark 4.x tarball into S3 if absent, and return its S3 URI.

    Downloads from the Apache archive (which retains all releases) on first use, then caches it in the SCT
    bucket so the EMR bootstrap can pull it from S3 - independent of Apache live-mirror rotation for latest release.

    Args:
        s3_bucket: Target S3 bucket name.
        region_name: AWS region for the S3 client.

    Returns:
        str: S3 URI of the Spark tarball.
    """
    s3_uri = f"s3://{s3_bucket}/{_SPARK4_S3_KEY}"
    s3_client = boto3.client("s3", region_name=region_name)
    _ensure_s3_bucket(s3_client, s3_bucket, region_name)

    try:
        s3_client.head_object(Bucket=s3_bucket, Key=_SPARK4_S3_KEY)
        LOGGER.info("Spark 4.x tarball already exists at %s, skipping upload", s3_uri)
        return s3_uri
    except ClientError as exc:
        if exc.response["Error"]["Code"] != "404":
            raise

    LOGGER.info("Downloading Spark 4.x tarball from %s", _SPARK4_ARCHIVE_URL)
    with tempfile.TemporaryDirectory() as tmp_dir:
        local_path = os.path.join(tmp_dir, f"{_SPARK4_PACKAGE}.tgz")
        response = requests.get(_SPARK4_ARCHIVE_URL, stream=True, timeout=600)
        response.raise_for_status()
        with open(local_path, "wb") as tarball_file:
            tarball_file.writelines(response.iter_content(chunk_size=8192))

        LOGGER.info("Uploading Spark 4.x tarball to %s", s3_uri)
        s3_client.upload_file(local_path, s3_bucket, _SPARK4_S3_KEY)

    LOGGER.info("Spark 4.x tarball uploaded to %s", s3_uri)
    return s3_uri


# TODO: Spark 4.x workaround
def build_spark4_bootstrap_actions(region_name, s3_bucket=None):
    """Mirror the Spark 4.x tarball to S3, upload the bootstrap script, and return EMR bootstrap actions."""
    s3_bucket = s3_bucket or f"sct-emr-spark-migrator-{region_name}"
    upload_spark4_tarball_to_s3(s3_bucket, region_name)
    s3_client = boto3.client("s3", region_name=region_name)
    s3_client.put_object(
        Bucket=s3_bucket, Key=_BOOTSTRAP_SCRIPT_S3_KEY, Body=_build_bootstrap_script(s3_bucket).encode("utf-8")
    )
    return [
        {
            "Name": f"Install Spark {SPARK4_VERSION}",
            "ScriptBootstrapAction": {"Path": f"s3://{s3_bucket}/{_BOOTSTRAP_SCRIPT_S3_KEY}"},
        },
    ]


def get_migrator_download_url(release_tag):
    """Build GitHub release asset URL for scylla-migrator JAR.

    Args:
        release_tag: GitHub release tag (e.g., 'v1.1.2') or full HTTPS URL.

    Returns:
        str: Download URL for the JAR asset.
    """
    if release_tag.startswith("https://"):
        return release_tag
    return f"https://github.com/{MIGRATOR_GITHUB_REPO}/releases/download/{release_tag}/{MIGRATOR_JAR_NAME}"


def upload_migrator_jar_to_s3(release_tag, s3_bucket, region_name, s3_prefix="jars"):
    """Download migrator JAR from GitHub and upload to S3 if not already present.

    Args:
        release_tag: GitHub release tag or direct URL.
        s3_bucket: Target S3 bucket name.
        region_name: AWS region for S3 client.
        s3_prefix: S3 key prefix (default: 'jars').

    Returns:
        str: S3 URI of the uploaded JAR (s3://bucket/prefix/tag/scylla-migrator-assembly.jar).
    """
    s3_key = f"{s3_prefix}/{release_tag}/{MIGRATOR_JAR_NAME}"
    s3_uri = f"s3://{s3_bucket}/{s3_key}"

    s3_client = boto3.client("s3", region_name=region_name)
    _ensure_s3_bucket(s3_client, s3_bucket, region_name)

    # check if the JAR already exists in S3
    try:
        s3_client.head_object(Bucket=s3_bucket, Key=s3_key)
        LOGGER.info("Migrator JAR already exists at %s, skipping upload", s3_uri)
        return s3_uri
    except ClientError as exc:
        if exc.response["Error"]["Code"] != "404":
            raise

    download_url = get_migrator_download_url(release_tag)
    LOGGER.info("Downloading migrator JAR from %s", download_url)

    with tempfile.TemporaryDirectory() as tmp_dir:
        local_path = os.path.join(tmp_dir, MIGRATOR_JAR_NAME)
        response = requests.get(download_url, stream=True, timeout=300)
        response.raise_for_status()
        with open(local_path, "wb") as jar_file:
            jar_file.writelines(response.iter_content(chunk_size=8192))

        LOGGER.info("Uploading migrator JAR to %s", s3_uri)
        s3_client.upload_file(local_path, s3_bucket, s3_key)

    LOGGER.info("Migrator JAR uploaded to %s", s3_uri)
    return s3_uri


def discover_cs_source_hosts(test_id: str, region_name: str) -> list[str]:
    """Return reachable IPs of running Cassandra source nodes tagged with test_id.

    Public IP is preferred (works cross-VPC). Falls back to private IP when no
    public IP is assigned (same-VPC / VPC-peered topology).

    Args:
        test_id: SCT cluster test_id for which cs-db EC2 instances should be discovered.
        region_name: AWS region to search in.

    Returns:
        List of IP addresses (public preferred, private fallback) for running cs-db instances.
    """
    ec2 = boto3.client("ec2", region_name=region_name)
    response = ec2.describe_instances(
        Filters=[
            {"Name": "tag:TestId", "Values": [test_id]},
            {"Name": "tag:NodeType", "Values": ["cs-db"]},
            {"Name": "instance-state-name", "Values": ["running"]},
        ]
    )
    hosts = [
        ip
        for reservation in response["Reservations"]
        for inst in reservation["Instances"]
        if (ip := inst.get("PublicIpAddress") or inst.get("PrivateIpAddress"))
    ]
    if not hosts:
        raise RuntimeError(f"no running cs-db instances found for test_id={test_id} in {region_name}")
    LOGGER.info("Discovered %d cs-db host(s) for test_id=%s in %s: %s", len(hosts), test_id, region_name, hosts)
    return hosts


@dataclass
class MigratorConfig:
    """Configuration for a spark-migrator job.

    Args:
        source_type: Source database type for scylla-migrator ('cassandra' or 'scylla').
        source_hosts: Source database hosts (one or more contact points).
        source_port: Source database CQL port.
        source_keyspace: Source keyspace name.
        source_table: Source table name.
        target_host: Target database host (Scylla endpoint).
        target_port: Target database CQL port.
        target_keyspace: Target keyspace name.
        target_table: Target table name.
        config_path: S3 path to migrator config YAML (alternative to individual fields).
        spark_executor_memory: Spark executor memory (e.g., '4g').
        spark_executor_cores: Number of Spark executor cores.
        spark_parallelism: Spark default parallelism.
        extra_spark_args: Additional spark-submit arguments.
    """

    source_type: str = "cassandra"
    source_hosts: list[str] = field(default_factory=list)
    source_port: int = 9042
    source_keyspace: str = ""
    source_table: str = ""
    target_host: str = ""
    target_port: int = 9042
    target_keyspace: str = ""
    target_table: str = ""
    config_path: str = ""
    spark_executor_memory: str = "4g"
    spark_executor_cores: int = 2
    spark_parallelism: int = 200
    extra_spark_args: list = field(default_factory=list)


class SparkMigratorRunner:
    """Submits and monitors spark-migrator jobs on EMR clusters.

    Args:
        emr_provisioner: An EmrClusterProvisioner instance managing the EMR cluster.
    """

    def __init__(self, emr_provisioner):
        self.emr_provisioner = emr_provisioner

    # TODO: Spark 4.x workaround
    @property
    def s3_bucket(self):
        return f"sct-emr-spark-migrator-{self.emr_provisioner.region_name}"

    def _upload_runner_script(
        self,
        jar_s3_path,
        config_s3_path,
        main_class=_MIGRATOR_MAIN_CLASS,
        script_key=_RUNNER_SCRIPT_S3_KEY,
        deploy_mode="cluster",
    ):
        """Generate and upload a wrapper script that invokes standalone Spark 4.x.

        Args:
            jar_s3_path: S3 path to the spark-migrator JAR.
            config_s3_path: S3 path to the migrator config YAML.
            main_class: Spark application main class (Migrator or Validator).
            script_key: S3 key for the uploaded wrapper script (distinct keys
                avoid migrator/validator scripts overwriting each other).
            deploy_mode: Spark deploy mode — "cluster" (default, driver on a
                YARN executor) or "client" (driver on the master). Validator
                uses "client" so its log4j output lands in the EMR step's
                stderr.gz on S3 instead of YARN aggregated container logs.
        """
        script_content = f"""#!/bin/bash
set -ex

# download the migrator JAR and config from S3
aws s3 cp {jar_s3_path} /tmp/scylla-migrator-assembly.jar
aws s3 cp {config_s3_path} /tmp/migrator-config.yaml

# point YARN at standalone Spark 4.0 (EMR's bundled Spark ships Scala 2.12,
# but the migrator requires Scala 2.13 which comes with Spark 4.0)
export SPARK_HOME={SPARK4_INSTALL_DIR}
export HADOOP_CONF_DIR=/etc/hadoop/conf
export YARN_CONF_DIR=/etc/hadoop/conf

cd /tmp

{SPARK4_INSTALL_DIR}/bin/spark-submit \\
  --deploy-mode {deploy_mode} \\
  --master yarn \\
  --class {main_class} \\
  --conf "spark.scylla.config=migrator-config.yaml" \\
  --files /tmp/migrator-config.yaml \\
  /tmp/scylla-migrator-assembly.jar
"""
        region_name = self.emr_provisioner.region_name
        s3_client = boto3.client("s3", region_name=region_name)
        _ensure_s3_bucket(s3_client, self.s3_bucket, region_name)
        s3_client.put_object(
            Bucket=self.s3_bucket,
            Key=script_key,
            Body=script_content.encode("utf-8"),
        )
        s3_uri = f"s3://{self.s3_bucket}/{script_key}"
        LOGGER.info("Runner script uploaded to %s", s3_uri)
        return s3_uri

    def _use_bootstrap(self):
        """True when the legacy Spark-4 bootstrap workaround is selected (emr_install_spark4_via_bootstrap)."""
        return bool(self.emr_provisioner.params.get("emr_install_spark4_via_bootstrap"))

    def _submit_spark_step(self, cluster_id, step_name, jar_s3_path, config_s3_path, main_class, deploy_mode):
        """Submit a Spark EMR step using either legacy bootstrap mode or native command-runner mode.

        Args:
            cluster_id: EMR cluster ID.
            step_name: EMR step name.
            jar_s3_path: S3 path to the spark-migrator JAR.
            config_s3_path: S3 path to the migrator config YAML.
            main_class: Spark application main class (Migrator or Validator).
            deploy_mode: Spark deploy mode ("cluster" for migration, "client" for the validator so its
                log4j output lands in the step's stderr.gz).

        Returns:
            str: Step ID of the submitted job.
        """
        if self._use_bootstrap():  # Spark 4.x workaround - legacy script-runner.jar + /opt/spark4 wrapper
            script_key = _RUNNER_SCRIPT_S3_KEY if main_class == _MIGRATOR_MAIN_CLASS else _VALIDATOR_SCRIPT_S3_KEY
            runner_script_uri = self._upload_runner_script(
                jar_s3_path,
                config_s3_path,
                main_class=main_class,
                script_key=script_key,
                deploy_mode=deploy_mode,
            )
            jar = f"s3://{self.emr_provisioner.region_name}.elasticmapreduce/libs/script-runner/script-runner.jar"
            args = [runner_script_uri]
            return self.emr_provisioner.add_step(cluster_id=cluster_id, step_name=step_name, jar=jar, args=args)

        # native EMR Spark 4
        jar = _COMMAND_RUNNER_JAR
        config_basename = config_s3_path.rsplit("/", 1)[-1]
        config_ref = config_basename
        files_arg = config_s3_path

        pre_cmd = ""
        if deploy_mode == "client":
            config_ref = f"/tmp/{config_basename}"
            files_arg = config_ref
            pre_cmd = f"aws s3 cp {config_s3_path} {config_ref} && "

        spark_submit_parts = [
            "spark-submit",
            "--deploy-mode",
            deploy_mode,
            "--master",
            "yarn",
            "--class",
            main_class,
            "--conf",
            f"{_SPARK_SCYLLA_CONFIG_KEY}={config_ref}",
            "--files",
            files_arg,
            jar_s3_path,
        ]

        if pre_cmd:
            args = ["bash", "-c", pre_cmd + " ".join(spark_submit_parts)]
        else:
            args = spark_submit_parts

        return self.emr_provisioner.add_step(cluster_id=cluster_id, step_name=step_name, jar=jar, args=args)

    def submit_migration_job(self, cluster_id, jar_s3_path, migrator_config):
        """Submit a spark-migrator job as an EMR step (cluster deploy mode).

        Args:
            cluster_id: EMR cluster ID.
            jar_s3_path: S3 path to the spark-migrator JAR.
            migrator_config: MigratorConfig with config_path set to S3 URI.

        Returns:
            str: Step ID of the submitted job.
        """
        step_id = self._submit_spark_step(
            cluster_id, "spark-migrator", jar_s3_path, migrator_config.config_path, _MIGRATOR_MAIN_CLASS, "cluster"
        )
        LOGGER.info("Spark migrator job submitted as step %s on cluster %s", step_id, cluster_id)
        return step_id

    def submit_validator_job(self, cluster_id, jar_s3_path, migrator_config):
        """Submit a scylla-migrator validator job as an EMR step (client deploy mode).

        Validator reads the same migrator config YAML as the migration — no separate config upload
        needed. Client mode keeps the driver on the master so its log4j diff summary lands in the
        step's stderr.gz.

        Args:
            cluster_id: EMR cluster ID.
            jar_s3_path: S3 path to the spark-migrator JAR (same as migration).
            migrator_config: MigratorConfig with config_path set to S3 URI.

        Returns:
            str: Step ID of the submitted validator job.
        """
        step_id = self._submit_spark_step(
            cluster_id,
            "spark-migrator-validator",
            jar_s3_path,
            migrator_config.config_path,
            _VALIDATOR_MAIN_CLASS,
            "client",
        )
        LOGGER.info("Spark migrator validator job submitted as step %s on cluster %s", step_id, cluster_id)
        return step_id

    def get_step_stdout(self, cluster_id, step_id):
        """Fetch step stdout + stderr from the EMR cluster's S3 log directory.

        EMR uploads each step's stdout.gz and stderr.gz under ``<LogUri>/<cluster_id>/steps/<step_id>/``
        after the step completes. Both streams are downloaded, gunzipped, decoded and concatenated
        so the validator parser can scan them as a single text blob (the validator emits its diff
        summary to stderr via log4j; stdout is generally empty).

        Args:
            cluster_id: EMR cluster ID.
            step_id: Step ID whose logs to fetch.

        Returns:
            str: Concatenated decoded text from stdout.gz + stderr.gz. Missing
            log files are skipped silently (returns "" if neither exists).
        """
        cluster_info = self.emr_provisioner.emr_client.describe_cluster(ClusterId=cluster_id)
        log_uri = cluster_info["Cluster"].get("LogUri", "").rstrip("/")
        # EMR normalises s3:// to s3n:// (legacy Hadoop S3 protocol); accept both
        # (plus s3a:// for completeness) so the prefix split works regardless of variant
        scheme_prefix = next((p for p in ("s3://", "s3n://", "s3a://") if log_uri.startswith(p)), None)
        if not scheme_prefix:
            raise RuntimeError(f"cluster {cluster_id} has no S3 LogUri configured — cannot fetch step logs")

        bucket, _, prefix = log_uri[len(scheme_prefix) :].partition("/")
        key_prefix = f"{prefix.rstrip('/')}/{cluster_id}/steps/{step_id}"
        s3_client = boto3.client("s3", region_name=self.emr_provisioner.region_name)

        text_parts = []
        for log_name in ("stdout.gz", "stderr.gz"):
            try:
                obj = s3_client.get_object(Bucket=bucket, Key=f"{key_prefix}/{log_name}")
                text_parts.append(gzip.decompress(obj["Body"].read()).decode("utf-8", errors="replace"))
            except ClientError as exc:
                if exc.response["Error"]["Code"] in ("NoSuchKey", "404"):
                    continue
                raise
        return "\n".join(text_parts)

    def wait_for_migration(self, cluster_id, step_id, timeout_minutes=360):
        """Wait for a migration/validator job to complete.

        Args:
            cluster_id: EMR cluster ID.
            step_id: Step ID of the migration job.
            timeout_minutes: Max wall-clock minutes. Default 360 (6h) covers
                multi-100GB Cassandra→Scylla migrations. Validator callers
                should pass a smaller value (~60).

        Returns:
            dict: Step status after completion.
        """
        return self.emr_provisioner.wait_for_step_completion(cluster_id, step_id, timeout_minutes=timeout_minutes)

    def run_migration(self, cluster_id, jar_s3_path, migrator_config, timeout_minutes=360):
        """Submit and wait for a migration job to complete.

        Args:
            cluster_id: EMR cluster ID.
            jar_s3_path: S3 path to the spark-migrator JAR.
            migrator_config: MigratorConfig with config_path set to S3 URI.
            timeout_minutes: Max wall-clock minutes for the migration step.

        Returns:
            dict: Migration result with step_id, status, and duration.
        """
        start_time = time.time()
        step_id = self.submit_migration_job(cluster_id, jar_s3_path, migrator_config)
        status = self.wait_for_migration(cluster_id, step_id, timeout_minutes=timeout_minutes)
        duration = time.time() - start_time

        LOGGER.info("Migration completed in %.1f seconds", duration)
        return {
            "step_id": step_id,
            "status": status,
            "duration_seconds": duration,
        }

    def get_step_logs(self, cluster_id, step_id):
        """Retrieve step logs from the EMR cluster.

        Args:
            cluster_id: EMR cluster ID.
            step_id: Step ID.

        Returns:
            dict: Step status information including any failure details.
        """
        return self.emr_provisioner.get_step_status(cluster_id, step_id)

    @staticmethod
    def generate_migrator_config(migrator_config):
        """Generate a scylla-migrator configuration dictionary from MigratorConfig.

        Args:
            migrator_config: MigratorConfig with source/target parameters.

        Returns:
            dict: Configuration dictionary matching scylla-migrator's expected YAML format.
        """
        return {
            "source": {
                "type": migrator_config.source_type,
                "host": migrator_config.source_hosts[0],
                "port": migrator_config.source_port,
                "keyspace": migrator_config.source_keyspace,
                "table": migrator_config.source_table,
                "consistencyLevel": "LOCAL_QUORUM",
                "preserveTimestamps": True,
                "splitCount": 256,
                "connections": 8,
                "fetchSize": 1000,
                "stripTrailingZerosForDecimals": False,
            },
            "target": {
                "type": "scylla",
                "host": migrator_config.target_host,
                "port": migrator_config.target_port,
                "keyspace": migrator_config.target_keyspace,
                "table": migrator_config.target_table,
                "consistencyLevel": "LOCAL_QUORUM",
                "connections": 16,
                "stripTrailingZerosForDecimals": False,
            },
            "savepoints": {
                "path": "/tmp/savepoints",
                "intervalSeconds": 300,
            },
            # required by com.scylladb.migrator.Validator (main class fails fast with
            # "Missing required property 'validation' in the configuration file" when
            # absent).
            "validation": {
                "compareTimestamps": False,
                "ttlToleranceMillis": 0,
                "writetimeToleranceMillis": 0,
                "failuresToFetch": 100,
                "floatingPointTolerance": 0.0,
                "timestampMsTolerance": 0,
            },
        }

    def upload_config_to_s3(self, config_dict, s3_bucket, test_id, region_name):
        """Upload a migrator config dictionary as YAML to S3.

        Args:
            config_dict: Configuration dictionary to serialize.
            s3_bucket: Target S3 bucket name.
            test_id: Test ID used to namespace the config file.
            region_name: AWS region for S3 client.

        Returns:
            str: S3 URI of the uploaded config file.
        """
        s3_key = f"configs/{test_id}/config.yaml"
        s3_uri = f"s3://{s3_bucket}/{s3_key}"

        s3_client = boto3.client("s3", region_name=region_name)
        _ensure_s3_bucket(s3_client, s3_bucket, region_name)
        s3_client.put_object(
            Bucket=s3_bucket,
            Key=s3_key,
            Body=yaml.dump(config_dict, default_flow_style=False).encode("utf-8"),
        )

        LOGGER.info("Migrator config uploaded to %s", s3_uri)
        return s3_uri
