#!/usr/bin/env python3

import argparse
import typing
import subprocess
import logging
import shutil
import shlex
import os

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("MaelstromInvoker")


DESCRIPTION = '''
Run the maelstrom tests for a given fly.io dist-sys challenge.

Currently, the following challenges are implemented:

- 01: Echo
- 02: Unique ID Generation
- 03a: Broadcast (Single Node)
- 03b: Broadcast (Multi-Node)

'''


IMPLEMENTED_CHALLENGES = {
    "1": "Echo", 
    "2": "Unique ID Generation", 
    "3a": "Broadcast (Single Node)",
    "3b": "Broadcast (Multi-Node)",
    # "3c": "Broadcast (Fault Tolerant)"
}

CHALLENGE_COMMANDS = {
    "1": "./maelstrom test -w echo --bin solutions/01-echo --time-limit 5",
    "2": "./maelstrom test -w unique-ids --bin solutions/02-unique-id-generation --time-limit 30 --rate 1000 --node-count 3 --availability total --nemesis partition",
    "3a": "./maelstrom test -w broadcast --bin solutions/03a-broadcast-single-node --node-count 1 --time-limit 20 --rate 10",
    "3b": "./maelstrom test -w broadcast --bin solutions/03b-broadcast-multi-node --node-count 5 --time-limit 20 --rate 10",
    # "3c": "./maelstrom test -w broadcast --bin solutions/03c-broadcast-fault-tolerant --node-count 5 --time-limit 20 --rate 10 --nemesis partition",
}

def parse_args():
    parser = argparse.ArgumentParser(prog="MaelstromInvoker", description=DESCRIPTION)
    parser.add_argument(
        "challenge",
        metavar="CHALLENGE_ID",
        nargs="+",
        type=str,
        help="The ids of the challenges to run.",
    )

    args = parser.parse_args()

    if any(
        challenge_id not in IMPLEMENTED_CHALLENGES for challenge_id in args.challenge
    ):
        raise ValueError(
            f"Invalid or unimplemented challenge id. Recognized challenge ids are: {str(IMPLEMENTED_CHALLENGES)}"
        )
    return args


def build_binaries():
    return subprocess.run(["cargo", "build", "--release"], check=True)


def copy_binaries_to_maelstrom_solutions():
    """Copy over the generated Rust binaries to maelstrom's solutions directory."""

    SRC_DIR = "target/release"
    DEST_DIR = "maelstrom/solutions"

    for file_path in os.listdir(SRC_DIR):
        src_file_path = os.path.join(SRC_DIR, file_path)
        # Copy over all the executables.

        if os.path.isfile(src_file_path) and os.access(src_file_path, os.X_OK):
            dest_file_path = os.path.join(DEST_DIR, file_path)
            shutil.copy(src_file_path, dest_file_path)
            logging.debug("Copied %s to %s", src_file_path, dest_file_path)


def run_test(challenge: str):
    """Run maelstrom for a given challenge."""

    command = CHALLENGE_COMMANDS.get(challenge, None)

    if command is None:
        raise ValueError(f"Couldn't find a test for challenge: {str(challenge)}")

    logger.info(
        'Running maelstrom test for challenge "%s" (id: %s)',
        IMPLEMENTED_CHALLENGES[challenge],
        challenge,
    )

    result = subprocess.run(shlex.split(command), check=True, cwd="maelstrom")
    return result


def main():
    args = parse_args()
    logger.info("Arguments: %s", args)
    build_binaries()
    copy_binaries_to_maelstrom_solutions()
    for challenge_id in args.challenge:
        run_test(challenge_id)


if __name__ == "__main__":
    main()