import epicbox
import argparse
from multiprocessing import Pool, Manager
from string import Template

import comparator
import languages

parser = argparse.ArgumentParser()

parser.add_argument("filename", help="Input file name.", type=str)
parser.add_argument("--inputs", type=str, nargs="+")
parser.add_argument("--outputs", type=str, nargs="+")
parser.add_argument(
    "--prog_language",
    help="Programming language of input.",
    type=str,
    default="python",
    choices=[*languages.supported.keys()],
)
parser.add_argument(
    "--time", help="Maximum time of program execution in ms.", type=int, default=1000
)
parser.add_argument(
    "--memory", help="Maximum RAM a program can use in mb.", type=int, default=250
)

args = parser.parse_args()

epicbox.configure(profiles=languages.supported[args.prog_language]["profiles"])

source = (
    "1" + args.filename[args.filename.rfind(".") :]
    if args.filename.rfind(".") != -1
    else "a"
)
language = languages.supported[args.prog_language]

exec = "a.out"
files = [{"name": source, "content": bytes(open(args.filename).read(), "utf-8")}]
limits = {"cputime": args.time, "memory": args.memory}

if not args.inputs or not args.outputs:
    exit(0)


def evaluate(exec, files, inStr, outStr, i, madDict):
    print("Started input: " + str(i))
    result = epicbox.run(
        language["run"]["docker_profile"],
        language["run"]["command"].substitute(exec=exec, source=source),
        limits=limits,
        files=files + [{"name": "in.txt", "content": bytes(inStr, "utf-8")}],
        stdin=inStr,
        workdir=madDict["work_dir"],
    )
    print(result)
    return comparator.default(result["stdout"].decode("utf-8"), outStr)


if __name__ == "__main__":
    with epicbox.working_directory() as workdir:
        if language["compile"] != {}:
            result = epicbox.run(
                language["compile"]["docker_profile"],
                language["compile"]["command"].substitute(exec=exec, source=source),
                files=files,
                workdir=workdir,
            )
        manager = Manager()

        d = manager.dict()
        d["work_dir"] = workdir

        with Pool(1) as p:
            evalData = [
                (
                    exec,
                    files if language["compile"] == {} else [],
                    open(i).read(),
                    open(o).read(),
                    j,
                    d,
                )
                for j, (i, o) in enumerate(zip(args.inputs, args.outputs))
            ]
            print(p.starmap(evaluate, evalData))