import argparse
import json
import math
import sys
import os
import subprocess

from zipfile import ZipFile

"""
    Checks if all submission output the correct decrypted data.
    Directory structure should be as follows:

    submissions
    ├── <team1>
    |   ├── mem_layout.json [optional]
    |   └── <savename>.zip
    ├── <team2>
    |   ├── mem_layout.json [optional]
    |   └── <savename>.zip
    ...

    group
    ├── sBox
    ├── key
    ├── data_encrypted
    └── data_decrypted

    mem_layout.json must contain the keys 'sbox', 'key' and 'data', with the respective addresses as values.
    If the file does not exist, the user will be prompted to enter these addresses. The data
    will then be saved as 'mem_layout.json', to be used in subsequent calls.

    If more than one simulator ZIP file is present, the user will also be prompted to choose.
"""

SIMULATOR_PATH = "./minimax_simulator-2.0.0-cli.jar"

# helper strings
OK = "\033[32mOK\033[0m"
ERROR = "\033[31mERROR\033[0m"


def compare_result(actual_file: str, expected_file: str) -> bool:
    """
    Compares the file at path 'actual_file' with the file at path 'expected_file' bytewise.
    If the actual file is larger than the expected file, additional bytes will be ignored.
    """
    with open(actual_file, "rb") as actual, open(expected_file, "rb") as expected:
        actual_bytes = actual.read()
        expected_bytes = expected.read()

    return actual_bytes[: len(expected_bytes)] == expected_bytes


def create_mem_layout() -> dict:
    """
    Prompts the user to enter addresses for the sbox, key and data.
    Returns a dictionary with these three keys.
    """
    mem_layout = dict()
    for key in ("sbox", "key", "data"):
        while True:
            try:
                value = input(f"{key} address (prefix hex numbers with '0x'):   ")
                base = 16 if value.startswith("0x") else 10
                value = int(value, base=base)
                break
            except ValueError:
                print("Adress is not an integer.")
            except KeyboardInterrupt:
                sys.exit("Aborted.")
        mem_layout[key] = value
    return mem_layout


def is_valid_zip(file: str) -> bool:
    """
    Checks if a zip file is a save file from the minimax simulator, e.g. if the contents
    are a 'machine.json' and 'signal.json' file.
    """
    if not file.endswith(".zip"):
        return False
    with ZipFile(file) as machine_zip:
        zip_content = machine_zip.namelist()
    return set(zip_content) == set(("machine.json", "signal.json"))


def select_zip(zips: list) -> str:
    """
    Prompts the user to select a single zip file from a list, and returns it.
    """
    print("Multiple zip files found. Please select one:")
    for index, f in enumerate(zips, start=1):
        print(f"[{index}] {f}")
    while True:
        try:
            selection = input("Enter the number of the zip file to select:    ")
            selection = int(selection) - 1
            if selection <= 0 or selection > len(zips):
                print(f"Please select a number between 1 and {len(zips)}.")
            else:
                return zips[selection]
        except ValueError:
            print("Please enter a valid integer.")
        except KeyboardInterrupt:
            sys.exit("Aborted")


def evaluate(
    zip_file: str,
    sbox_file: str,
    key_file: str,
    data_file: str,
    result_file: str,
    mem_layout: dict,
    simulator: str,
) -> None:
    """
    Runs the minimax simulator on the given input. The resulting file is saved in 'result_file'.
    """
    args = [
        "java",
        "-jar",
        simulator,
        zip_file,
        "--import-file",
        sbox_file,
        "--import-from",
        mem_layout["sbox"],
        "--import-file",
        key_file,
        "--import-from",
        mem_layout["key"],
        "--import-file",
        data_file,
        "--import-from",
        mem_layout["data"],
        "--export-file",
        result_file,
        "--export-from",
        mem_layout["data"],
        "--export-to",
        mem_layout["data"] + math.ceil(os.path.getsize(data_file)),
    ]

    args = [
        str(arg) for arg in args
    ]  # subprocess.run requires all arguments to be strings

    print(f"Running simulator, storing result in '{result_file}'")

    result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    print("\033[38;5;245m")

    print(result.stdout.decode("utf-8"))

    print("\033[38;5;124")

    print(result.stderr.decode("utf-8"))

    print("\033[0m")


if __name__ == "__main__":  # only run if executed as script

    parser = argparse.ArgumentParser()
    parser.add_argument("submissions", type=str, help="Submissions root directory")
    parser.add_argument(
        "group", type=str, help="Group directory, contains all project files"
    )
    parser.add_argument(
        "-e",
        "--file-extension",
        dest="file_ext",
        type=str,
        help="Result file extension",
    )
    parser.add_argument("-j", "--jar", dest="jar", type=str, help="Simulator jar file")
    args = parser.parse_args()

    # Load teams
    teams = [
        e
        for e in os.listdir(args.submissions)
        if os.path.isdir(os.path.join(args.submissions, e))
    ]

    print(f"The following teams were found:")
    for team in teams:
        print(f"* {team}")

    # Check directory structure
    for file in ("sBox", "key", "data_encrypted", "data_decrypted"):
        if not os.path.exists(os.path.join(args.group, file)):
            sys.exit(f"Group project file '{file}' is missing.")

    # Check file extension
    if args.file_ext is None:
        args.file_ext = ""
    elif not args.file_ext.startswith("."):
        args.file_ext = f".{args.file_ext}"

    # Update simulator path if given
    simulator = SIMULATOR_PATH if args.jar is None else args.jar

    # Store evaluation results
    expected_file = os.path.join(args.group, "data_decrypted")
    results = []

    # Evaluate each team
    for team in teams:

        # load memory layout file if available (otherwise create and store it)
        try:
            with open(
                os.path.join(args.submissions, team, "mem_layout.json"), "r"
            ) as mem_layout_file:
                mem_layout = json.load(mem_layout_file)
        except FileNotFoundError:
            mem_layout = create_mem_layout()
            with open(
                os.path.join(args.submissions, team, "mem_layout.json"), "w"
            ) as mem_layout_file:
                json.dump(mem_layout, mem_layout_file)

        # Select project file (if more there is more than one zip file)
        zip_files = [
            os.path.join(args.submissions, team, f)
            for f in os.listdir(os.path.join(args.submissions, team))
            if is_valid_zip(os.path.join(args.submissions, team, f))
        ]
        zip_file = zip_files[0] if len(zip_files) == 1 else select_zip(zip_files)

        # check memory layout and convert hexadecimal addresses
        for key in ("sbox", "key", "data"):
            if key not in mem_layout:
                sys.exit(f"memory_layout does not contain key '{key}'")
            try:
                addr = int(mem_layout[key])
            except ValueError:
                try:
                    addr = int(mem_layout[key], base=16)
                except ValueError:
                    sys.exit(f"Invalid address '{mem_layout[key]}' for key '{key}'")
            mem_layout[key] = addr

        # evaluate team
        sbox_file = os.path.join(args.group, "sBox")
        key_file = os.path.join(args.group, "key")
        data_file = os.path.join(args.group, "data_encrypted")
        result_file = os.path.join(
            args.submissions, team, f"data_decrypted{args.file_ext}"
        )

        evaluate(
            zip_file,
            sbox_file,
            key_file,
            data_file,
            result_file,
            mem_layout,
            simulator=simulator,
        )

        # Compare against expected result
        results.append(compare_result(result_file, expected_file))

    # Print result summary
    print("Summary:")
    for team, result in zip(teams, results):
        if result is True:
            print(f"[{OK}]    - {team}")
        else:
            print(f"[{ERROR}] - {team}")

    print()