#!/usr/bin/python3

# This script generates shell script that sets the the docker "--cpuset-cpus="
# option for a containers in the start.sh script.
#
# It assumes that this will only be used to generate assignments for a single
# host machine, or if bots run in multiple places, the same assignments can be
# used on all hosts.
#
# Change the "allocations" dictionary below to change core allocations,
# the actual numbering will be created and validated by this script.

import argparse
import logging
import sys

# Not really needed but makes it easier to read when debugging.
from collections import OrderedDict
from pprint import pformat

# container name suffix -> desired number of cores
allocations = OrderedDict([
    # Libcxx has to fit within a 2 hour time limit.
    ("-armv8-libcxx-01", 12),
    ("-armv8-libcxx-02", 12),
    ("-armv8-libcxx-03", 12),
    ("-armv8-libcxx-04", 12),
    ("-aarch64-libcxx-01", 10),
    ("-aarch64-libcxx-02", 10),
    ("-aarch64-libcxx-03", 10),
    ("-aarch64-libcxx-04", 10),
    ("-aarch64-libcxx-05", 10),
    # 2 stage bots, 15 cores each.
    ("-armv8-lld-2stage", 15),
    ("-armv7-vfpv3-2stage", 15),
    ("-armv7-2stage", 15),
    # 2 stage aarch64 needs more for unknown reasons.
    ("-aarch64-lld-2stage", 20),
    # Global Isel bots, 15 cores each due to instability in Clangd
    # tests without fixed cores.
    ("-aarch64-global-isel", 15),
    ("-armv7-global-isel", 15),
    # GCC building Flang uses upwards of 10GB per core.
    ("-flang-aarch64-latest-gcc", 16),
    # Debug builds can be even worse.
    ("-flang-aarch64-debug-reverse-iteration", 16),
    # Some LLDB tests are scheduling sensitive.
    ("-lldb-arm-ubuntu", 16),
])

# Amp-01 has 256 cores
MAX_CORES = 256

# Generate actual layout from the desired allocations.
def generate_layout(allocation):
    # container name suffix -> begin, end core
    layout = OrderedDict()
    current_core = 0
    used = 0
    for container, cores in allocation.items():
        if (current_core >= MAX_CORES) or ((current_core + cores - 1) >= MAX_CORES):
            raise RuntimeError(
                f"Core layout for is not valid, not enough cores!"
            )

        if container in layout.keys():
            raise RuntimeError(
                f"Container name suffix {container} is in layout more than once!"
            )

        layout[container] = (current_core, current_core + cores - 1)
        current_core += cores
        used += cores

    logging.debug(f"Layout:")
    logging.debug(pformat(layout))

    logging.debug(
        f"Uses {used} of {MAX_CORES} available cores, {MAX_CORES-used} remain."
    )

    return layout


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()

    if args.debug:
        logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)

    logging.debug("Desired allocations:")
    logging.debug(pformat(allocations))

    layout = generate_layout(allocations)
    print("  ### Auto generated by generate_core_assignments.py")
    print('  case "$botname" in')
    for container_suffix, cores in layout.items():
        start_core, end_core = cores
        print(f'    *{container_suffix}) cpuset_cpus="--cpuset-cpus={start_core}-{end_core}" ;;')
    print('  esac')
    print("  ### Auto generated code ends")