import argparse
import os
import subprocess

parser = argparse.ArgumentParser()
options = None

# unit: MB/s
instance2baseline = {
    # x864 series
    "m5.large": 81.25,
    "c5.xlarge": 143.75,
    "m5.xlarge": 143.75,
    "c5.2xlarge": 287.5,
    "m5.2xlare": 287.5,
    "r5.2xlarge": 287.5,
    "c5.4xlarge": 593.75,
    "m5.4xlarge": 593.75,
    "r5.4xlarge": 593.75,

    # arm series
    # the baseline throughput of m6g.large is too small, we shouldn't use the baseline
    # "m6g.large": 78.75,
    "c6g.xlarge": 148.5,
    "m6g.xlarge": 148.5,
    "c6g.2xlarge": 296.88,
    "m6g.2xlare": 296.88,
    "r6g.2xlarge": 296.88,
    "c6g.4xlarge": 593.75,
    "m6g.4xlarge": 593.75,
    "r6g.4xlarge": 593.75,
}


def is_arm_instance(instance_name: str) -> bool:
    return "g." in instance_name


def revise_tikv(tikv_path: str, instance_name: str, instance_throughput: float):
    io_rate_limit_bytes = int(instance_throughput * 1024 * 1024)

    cmd_str = f'''
    gsed -ri 's/max-bytes-per-sec\s*=.*$/max-bytes-per-sec = {io_rate_limit_bytes} # The baseline throughput of {instance_name} with EBS is {instance_throughput}MB\/s/g' \
    {tikv_path}'''
    subprocess.run(cmd_str, shell=True)


def revise_tiflash(tiflash_path: str, instance_name: str, instance_throughput: float):
    io_rate_limit_bytes = int(instance_throughput * 1024 * 1024)

    cmd_str = f'''
    gsed -ri 's/max_bytes_per_sec\s*=.*$/max_bytes_per_sec = {io_rate_limit_bytes} # The baseline throughput of {instance_name} with EBS is {instance_throughput}MB\/s./g' \
    {tiflash_path}'''
    subprocess.run(cmd_str, shell=True)


def revise_one_instance(instance_path: str, instance_name: str, instance_throughput: float):
    for root, dirs, files in os.walk(instance_path):
        for file in files:
            if file in ["tikv.toml"]:
                revise_tikv(os.path.join(root, file), instance_name, instance_throughput)
            if file in ["tiflash.toml"]:
                revise_tiflash(os.path.join(root, file), instance_name, instance_throughput)


def revise_one_version(version_path: str):
    global options

    instances_path = os.path.join(version_path, "instances")

    for root, dirs, files in os.walk(instances_path):
        for dirname in dirs:
            if not dirname.startswith("aws."):
                continue

            instance_name = dirname.removeprefix("aws.")

            if is_arm_instance(instance_name) and options.skip_arm:
                continue

            if not is_arm_instance(instance_name) and options.skip_x86:
                continue

            throughput = instance2baseline.get(instance_name)
            if throughput is None:
                continue

            revise_one_instance(os.path.join(root, dirname), instance_name, throughput)


def revise_all(versions_dir: str):
    for root, dirs, files in os.walk(versions_dir):
        for dirname in dirs:
            if dirname.startswith("v"):
                version_path = os.path.join(root, dirname)
                revise_one_version(version_path)


def run():
    global options
    parser.add_argument(
        "--versions-dir",
        dest="versions_dir",
        type=str,
        default="",
        help="the directory of versions",
    )
    parser.add_argument(
        "--skip-x86",
        dest="skip_x86",
        type=bool,
        default=False,
        help="skip revise throughput for X86 instances",
    )
    parser.add_argument(
        "--skip-arm",
        dest="skip_arm",
        type=bool,
        default=False,
        help="skip revise throughput for ARM instances",
    )
    options = parser.parse_args()

    revise_all(options.versions_dir)


if __name__ == '__main__':
    run()