#!/usr/bin/python3

# This script generates the docker "--cpuset-cpus=" option for a
# container with the given name.
#
# 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

# Machine name -> (container name suffix -> desired number of cores)
allocations = OrderedDict(
    (
        (
            "jade-01",
            # 2 stage bots, 15 cores each.
            OrderedDict(
                (
                    ("-armv8-lld-2stage", 15),
                    ("-aarch64-lld-2stage", 15),
                    ("-armv7-vfpv3-2stage", 15),
                    ("-armv7-2stage", 15),
                    # Global Isel bots, 15 cores each due to instability in Clangd
                    # tests without fixed cores.
                    ("-aarch64-global-isel", 15),
                    ("-armv7-global-isel", 15),
                )
            ),
        ),
        (
            "jade-04",
            # All libcxx containers get 8 cores.
            OrderedDict(
                (
                    ("-armv8-libcxx-01", 8),
                    ("-armv8-libcxx-02", 8),
                    ("-armv8-libcxx-03", 8),
                    ("-armv8-libcxx-04", 8),
                    ("-armv8-libcxx-05", 8),
                    ("-aarch64-libcxx-01", 8),
                    ("-aarch64-libcxx-02", 8),
                    ("-aarch64-libcxx-03", 8),
                    ("-aarch64-libcxx-04", 8),
                    ("-aarch64-libcxx-05", 8),
                    # This builds flang, give it 30 cores.
                    ("-aarch64-full-2stage", 30),
                )
            ),
        ),
    )
)

# This script only supports Jade machines at this time.
MAX_CORES = 160

# Generate actual layout from the desired allocations.
def generate_layout(machine, alloaction):

    # Machine name ending -> begin, end core
    layout = OrderedDict()
    current_core = 0
    used = 0
    for container, cores in allocation.items():
        if current_core >= MAX_CORES:
            raise RuntimeError(
                f"Core layout for machine {machine} 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 for machine {machine}:")
    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("containername")
    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))

    layouts = OrderedDict()
    for machine, allocation in allocations.items():
        layouts[machine] = generate_layout(machine, allocation)

    for machine, layout in layouts.items():
        for container_name_suffix, core_range in layout.items():
            # Only check suffix because we want to treat "normal-" and "silent-" containers
            # as the same. We assume we'll never run both on the same machine.
            if args.containername.endswith(container_name_suffix):
                logging.debug(
                    f"Found container {args.containername} on machine {machine}."
                )
                print(f"--cpuset-cpus={core_range[0]}-{core_range[1]}")
                exit(0)

    # If we get here, we have no specific allocation for the container.
    # This is expected.
    logging.debug(f"Did not find core allocation for container {args.containername}.")
    exit(0)