import argparse
import os
import subprocess
import tarfile


def parse_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument("-o", "--ort_master_path", required=True, help="ORT master repo")
    parser.add_argument("-t", "--tensorrt_home", required=True, help="TensorRT home directory")
    parser.add_argument("-c", "--cuda_home", required=True, help="CUDA home directory")
    parser.add_argument("-b", "--branch", required=False, default="master", help="Github branch to test perf off of")
    parser.add_argument("-s", "--save", required=False, help="Directory to archive wheel file")
    parser.add_argument("-a", "--use_archived", required=False, help="Archived wheel file")
    args = parser.parse_args()
    return args


def archive_wheel_file(save_path, ort_wheel_file):
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    subprocess.run(["cp", ort_wheel_file, save_path], check=True)


def install_new_ort_wheel(ort_master_path):
    ort_wheel_path = os.path.join(ort_master_path, "build", "Linux", "Release", "dist")
    p1 = subprocess.run(["find", ort_wheel_path, "-name", "*.whl"], stdout=subprocess.PIPE, check=True)
    stdout = p1.stdout.decode("utf-8").strip()
    ort_wheel = stdout.split("\n")[0]
    subprocess.run(["python3", "-m", "pip", "install", "--force-reinstall", ort_wheel], check=True)
    return ort_wheel


def main():
    args = parse_arguments()

    cmake_tar = "cmake-3.18.4-Linux-x86_64.tar.gz"
    if not os.path.exists(cmake_tar):
        subprocess.run(["wget", "-c", "https://cmake.org/files/v3.18/" + cmake_tar], check=True)
    tar = tarfile.open(cmake_tar)
    tar.extractall()
    tar.close()

    os.environ["PATH"] = os.path.join(os.path.abspath("cmake-3.18.4-Linux-x86_64"), "bin") + ":" + os.environ["PATH"]
    os.environ["CUDACXX"] = os.path.join(args.cuda_home, "bin", "nvcc")

    ort_master_path = args.ort_master_path
    pwd = os.getcwd()
    os.chdir(ort_master_path)

    if args.use_archived:
        ort_wheel_file = args.use_archived
        subprocess.run(["python3", "-m", "pip", "install", "--force-reinstall", ort_wheel_file], check=True)

    else:
        subprocess.run(["git", "fetch"], check=True)
        subprocess.run(["git", "checkout", args.branch], check=True)
        subprocess.run(["git", "pull", "origin", args.branch], check=True)
        subprocess.run(
            [
                "./build.sh",
                "--config",
                "Release",
                "--use_tensorrt",
                "--tensorrt_home",
                args.tensorrt_home,
                "--cuda_home",
                args.cuda_home,
                "--cudnn",
                "/usr/lib/x86_64-linux-gnu",
                "--build_wheel",
                "--skip_tests",
                "--parallel",
            ],
            check=True,
        )

        ort_wheel_file = install_new_ort_wheel(ort_master_path)

        if args.save:
            archive_wheel_file(args.save, ort_wheel_file)

    os.chdir(pwd)


if __name__ == "__main__":
    main()
