#!/usr/bin/env python3

import os
import sys
import shutil
import random
import subprocess

def build(algo_dir, template_dir="templates/linux"):
    build_dir = None
    while build_dir is None:
        r = "%09d" % random.randint(0, 999999999)
        d = os.path.join("build", r)
        if not os.path.isdir(d):
            build_dir = d
    
    shutil.copytree(algo_dir, build_dir)

    c = os.path.join(build_dir, "genkat_aead.c")
    if os.path.exists(c):
        os.remove(c)

    hfiles = []
    cfiles = []
    for r, d, f in os.walk(build_dir):
        for file in f:
            if file.endswith(".c"):
                cfiles.append(file)
            elif file.endswith(".h"):
                hfiles.append(file)

    for f in os.listdir(template_dir):
        dst = os.path.join(build_dir, f)
        src = os.path.join(template_dir, f)
        if os.path.isfile(src) or os.path.islink(src):
            shutil.copy2(src, dst)
        elif os.path.isdir(src):
            shutil.copytree(src, dst)
        else:
            raise Exception("I don't know what %s is" % src)

    env = os.environ
    env['SRC_FILES'] = ' '.join(cfiles)
    env['HDR_FILES'] = ' '.join(hfiles)

    wd = os.getcwd()
    try:
        os.chdir(build_dir)
        if os.path.isfile('./configure'):
            p = subprocess.Popen(["./configure"])
            p.wait()
            assert p.returncode == 0

        pargs = ['make']
        p = subprocess.Popen(['make'])
        p.wait()
        assert p.returncode == 0

    finally:
        os.chdir(wd)

    return build_dir

def find_test_vectors(d):
    kat = None
    while kat is None:
        if d == '':
            raise Exception("Test vector not found")
        for f in os.listdir(d):
            if f.startswith("LWC_AEAD_KAT_") and f.endswith(".txt"):
                if kat is not None:
                    raise Exception("Multiple test vectors?")
                kat = f
        d = os.path.split(d)[0]
    kat = os.path.join(d, kat)
    return kat

def main(argv):
    submissions_dir = "all-lwc-submission-files"
    template_dir = "templates/linux"
    if len(argv) > 1:
        template_dir = argv[1]
    print("Using template %s" % template_dir)
    subs = os.listdir(submissions_dir)

    # get all the submissions by looking for files named "encrypt.c"
    files = []
    for submission in subs:
        implementations_dir = os.path.join(submissions_dir, submission, "Implementations", "crypto_aead")

        if not os.path.isdir(implementations_dir):
            continue

        if "NOT ACCEPTED" in implementations_dir:
            continue

        print()
        print("###  %s  ###" % submission)

        c = 0
        # r=root, d=directories, f = files
        for r, d, f in os.walk(implementations_dir):
            for file in f:
                if file == "api.h":
                    f = os.path.join(r, file)
                    d = os.path.split(f)[0]
                    assert os.path.isdir(d)
                    print(d)
                    t = find_test_vectors(d)
                    print(t)
                    files.append(d)
                    c += 1

        if c == 0:
            raise Exception("No implementations found")



    # For testing, we only do the first
    files = files[:1]

    if os.path.isdir('build'):
        shutil.rmtree('build')

    print()

    for d in files:
        print()
        print(d)
        b = build(d)


if __name__ == "__main__":
    sys.exit(main(sys.argv))