import argparse
import json
import math
import os
import subprocess
import sys
from pathlib import Path

from util import is_valid_zip, select_zip

"""
    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 = Path("./minimax_simulator-2.0.0-cli.jar")

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


def ul(s: str) -> str:
    """
    Adds ansi escape sequences to underline string.
    """
    return f"\033[4m{s}\033[0m"


def compare_result(actual_file: Path, expected_file: Path) -> 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.
    """
    try:
        with open(actual_file, "rb") as actual, open(expected_file, "rb") as expected:
            actual_bytes = actual.read()
            expected_bytes = expected.read()
    except FileNotFoundError:
        return False

    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 evaluate(
    zip_file: Path,
    sbox_file: Path,
    key_file: Path,
    data_file: Path,
    result_file: Path,
    mem_layout: dict,
    simulator: Path,
) -> 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.get("result", mem_layout["data"]),
        "--export-to",
        mem_layout.get("result", mem_layout["data"]) + math.ceil(os.path.getsize(data_file) / 4),
    ]

    if "result" in mem_layout:
        print("Decryption was not done in-place.")

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

    print(f"Running simulator, storing result in '{result_file}'")
    print("\033[38;5;245m")

    subprocess.run(args, stdout=sys.stdout, stderr=sys.stderr)

    print("\033[0m")


if __name__ == "__main__":  # only run if executed as script
    parser = argparse.ArgumentParser(
        description="Runs the simulator on all projects found in the submissions directory."
    )
    parser.add_argument(
        "submissions",
        type=Path,
        help="Submissions root directory",
    )
    parser.add_argument(
        "group",
        type=Path,
        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=Path, help="Simulator jar file")
    args = parser.parse_args()

    # Load teams
    teams = [e for e in args.submissions.iterdir() if e.is_dir()]

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

    # Check directory structure
    for filename in ("sBox", "key", "data_encrypted", "data_decrypted"):
        if not (args.group / filename).exists():
            sys.exit(f"Group project file '{filename}' 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 = args.group / "data_decrypted"
    results = []

    # Evaluate each team
    for team in teams:
        print(ul(f"Evaluating team '{team.name}'"))

        # load memory layout file if available (otherwise create and store it)
        try:
            with open(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(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 = [file for file in team.glob("*.zip") if is_valid_zip(file)]
        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 = args.group / "sBox"
        key_file = args.group / "key"
        data_file = args.group / "data_encrypted"
        result_file = 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.name}")
        else:
            print(f"[{ERROR}] - {team.name}")

    print()