"""
Given a file containing the list of benchmarks, this script downloads the
benchmark files from the benchmark repository
and save them in the specified directory.

Usage: python download_benchmark.py --blist <file> --odir <directory>
Default values: --blist blist.txt --odir downloaded_benchmarks
"""

import os
import wget
import argparse

def download_benchmark_driver(benchmark_name, output_dir):
    DRV_PREFIX = "linux-4.2-rc1.tar.xz-32_7a-drivers--"
    DRV_SUFFIX = ".ko-entry_point.cil.out.i"

    BENCHMARK_REPO = "https://gitlab.com/sosy-lab/benchmarking/sv-benchmarks/-/raw/svcomp24/c/ldv-linux-4.2-rc1/"

    url = BENCHMARK_REPO + DRV_PREFIX + benchmark_name + DRV_SUFFIX

    output_file = os.path.join(output_dir, DRV_PREFIX + benchmark_name + DRV_SUFFIX)

    # Download the benchmark
    try:
        if os.path.exists(output_file):
            os.remove(output_file)
        wget.download(url, output_file)
    except:
        BENCHMARK_REPO_TODO = "https://gitlab.com/sosy-lab/benchmarking/sv-benchmarks/-/raw/svcomp24/c/ldv-linux-4.2-rc1-todo/"
        url = BENCHMARK_REPO_TODO + DRV_PREFIX + benchmark_name + DRV_SUFFIX
        output_file = os.path.join(output_dir, DRV_PREFIX + benchmark_name + DRV_SUFFIX)

        try:
            if os.path.exists(output_file):
                os.remove(output_file)
            wget.download(url, output_file)
        except:
            DRV_PREFIX_TMP = "linux-4.2-rc1.tar.xz-43_2a-drivers--"
            url = BENCHMARK_REPO + DRV_PREFIX_TMP + benchmark_name + DRV_SUFFIX
            output_file = os.path.join(output_dir, DRV_PREFIX_TMP + benchmark_name + DRV_SUFFIX)
            try:
                if os.path.exists(output_file):
                    os.remove(output_file)
                wget.download(url, output_file)
            except:
                print("Error: Benchmark {} not found".format(benchmark_name))

def download_benchmark_wcet(benchmark_name, output_dir):
    BENCHMARK_REPO = "https://gricad-gitlab.univ-grenoble-alpes.fr/pagai/pagai/-/raw/wcet_pagai_integration/benchmarks/wcet_base/"

    url = BENCHMARK_REPO + benchmark_name

    output_file = os.path.join(output_dir, benchmark_name)

    try:
        if os.path.exists(output_file):
            os.remove(output_file)
        wget.download(url, output_file)
    except:
        print("Error: Benchmark {} not found".format(benchmark_name))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--blist", type=str, help="File containing the list of benchmarks", default="wcet_list.txt")
    parser.add_argument("--odir", type=str, help="Output directory", default="downloaded_benchmarks")
    args = parser.parse_args()

    if not os.path.exists(args.odir):
        os.makedirs(args.odir)

    if not os.path.exists(args.blist):
        raise Exception("Error: File not found: " + args.blist)

    with open(args.blist, "r") as file:
        benchmarks = file.readlines()
    benchmarks = [benchmark.strip() for benchmark in benchmarks]

    if benchmarks[0].endswith(".c"):
        args.odir = os.path.join(args.odir, "wcet")
        os.makedirs(args.odir, exist_ok=True)
        for b in benchmarks:
            download_benchmark_wcet(b, args.odir)
    else:
        args.odir = os.path.join(args.odir, "drivers")
        os.makedirs(args.odir, exist_ok=True)
        for b in benchmarks:
            download_benchmark_driver(b, args.odir)
    print("\n{} benchmarks downloaded in {}".format(len(benchmarks), args.odir))

if __name__ == "__main__":
    main()