#!/usr/bin/python3

import os
import sys
import argparse
import subprocess
import json
import datetime
import random
import re
from dataclasses import dataclass


@dataclass
class Config:
    bench_directory: str
    output_directory: str
    compile_command: str
    run_command: str
    file_ext: str
    num_runs: int
    use_internal_timing: bool = False
    time_format: str = "gcomp"
    output_file: str = ""

    def __post_init__(self):
        if self.output_file == "":
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
            output_file = f"results-{timestamp}.json"
            output_file_path = os.path.join(self.output_directory, output_file)
            self.output_file = output_file_path


def find_benchmarks(config: Config):
    print("Finding benchmarks in directory: " + config.bench_directory)
    benchmarks = []
    with os.scandir(config.bench_directory) as entries:
        for entry in entries:
            if entry.is_file() and entry.name.endswith(config.file_ext):
                benchmarks.append(entry.name)
    print(f"Found {len(benchmarks)} benchmarks")
    return benchmarks


def compile_benchmarks(config: Config, benchmarks):
    for benchmark in benchmarks:
        print("Compiling benchmark: " + benchmark)
        source = os.path.join(config.bench_directory, benchmark)
        basename = os.path.splitext(benchmark)[0]
        output = os.path.join(config.output_directory, basename)
        cmd = config.compile_command.replace("{output}", output)
        cmd = cmd.replace("{source}", source)
        print(cmd)
        process = subprocess.run(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
        )
        stdout = process.stdout.decode("utf-8")
        stderr = process.stderr.decode("utf-8")
        print(f"stdout:\n{stdout}")
        print(f"stderr:\n{stderr}")
        if process.returncode != 0:
            print("Failed to compile benchmark: " + benchmark)
            sys.exit(1)


def run_benchmarks(config: Config, benchmarks):
    results = []
    for n in range(config.num_runs):
        print(f"Run {n + 1}/{config.num_runs}")
        random.shuffle(benchmarks)
        for benchmark in benchmarks:
            basename = os.path.splitext(benchmark)[0]
            print(f"Running benchmark: {basename}")
            path = os.path.join(config.output_directory, basename)
            cmd = f"{path} {config.run_command}"
            print(f"run: {cmd}")
            process = subprocess.run(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
            )
            stdout = process.stdout.decode("utf-8")
            stderr = process.stderr.decode("utf-8")
            print(f"stdout: {stdout}")
            print(f"stderr: {stderr}")
            if process.returncode != 0:
                print(f"Failed to run benchmark: {basename}")
                sys.exit(1)
            results.append(
                {
                    "basename": basename,
                    "run": n,
                    "stdout": stdout,
                    "stderr": stderr,
                }
            )
    return results


def compile_and_run_benchmarks(config: Config):
    # Create output directory if it does not exist
    os.makedirs(config.output_directory, exist_ok=True)
    benchmarks = find_benchmarks(config)
    print("Running benchmarks in directory: " + config.bench_directory)
    compile_benchmarks(config, benchmarks)
    print(f"Stack size (ulimit): {subprocess.getoutput('ulimit -s')}")
    print("All benchmarks compiled successfully")
    print("Running benchmarks in directory: " + config.output_directory)
    results = run_benchmarks(config, benchmarks)
    print("Done")
    return results


def write_results(config: Config, results):
    path = config.output_file
    output_dir = os.path.dirname(path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    with open(path, "w") as file:
        json.dump(results, file, indent=4)
    print(f"Results written to: {path}")


def parse_multitime_output(output: str):
    realtime_pattern = (
        r"real\s+(\d+\.\d+)\s+(\d+\.\d+)\s+\d+\.\d+\s+\d+\.\d+\s+\d+\.\d+"
    )
    match_opt = re.search(realtime_pattern, output)
    if match_opt is None:
        print("Failed to parse multitime output")
        sys.exit(1)
    mean = float(match_opt.group(1))
    # We ignore stdev, because we prefer using the disaggregated results
    return mean


def parse_internal_timing_output_gcomp(output: str) -> float:
    # Internal timing is printed using the following code in C:
    # `fprintf(stderr, "Execution time: %.6f seconds\n", elapsed);`
    timing_pattern = r"Execution time:\s+(\d+\.\d+)\s+seconds"
    match_opt = re.search(timing_pattern, output)
    if match_opt is None:
        print("Failed to parse internal timing output (expected gcomop format)")
        sys.exit(1)
    mean = float(match_opt.group(1))
    return mean


def parse_internal_timing_output_grift(output: str) -> float:
    # Follows this format:
    # time (sec): 0.386053
    timing_pattern = r"time \(sec\):\s*(\d+\.\d+)"
    match_opt = re.search(timing_pattern, output)
    if match_opt is None:
        print("Failed to parse internal timing output (expected grift format)")
        sys.exit(1)
    mean = float(match_opt.group(1))
    return mean


def parse_results(config: Config, results: list[dict]):
    if not config.use_internal_timing:
        parse = lambda result: parse_multitime_output(result["stderr"])
    elif config.time_format == "gcomp":
        parse = lambda result: parse_internal_timing_output_gcomp(result["stderr"])
    elif config.time_format == "grift":
        parse = lambda result: parse_internal_timing_output_grift(result["stdout"])
    else:
        raise ValueError(f"unexpected time format '{config.time_format}'")
    for result in results:
        runtime = parse(result)
        if runtime is None:
            raise ValueError(f"failed to parse time format")
        result["runtime"] = runtime
        del result["stderr"]
        del result["stdout"]
    return results


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "benchmarks_directory", help="Directory containing benchmark files"
    )
    parser.add_argument(
        "--compile-cmd",
        required=True,
        help="Command to compile the benchmarks. The occurrences of '{source}' and '{output}' will be replaced accordingly",
    )
    parser.add_argument(
        "--run-cmd", required=True, help="Command to run the benchmarks"
    )
    parser.add_argument(
        "--file-ext", default=".txt", help="File extension for source files"
    )
    parser.add_argument(
        "--num-runs", type=int, default=1, help="Number of runs for each sample"
    )
    parser.add_argument(
        "--internal-timing",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Use internal timing for benchmarks",
    )
    parser.add_argument(
        "--output-directory",
        default="output",
        help="Directory to store the compiled benchmarks and results",
    )
    parser.add_argument(
        "--output-file",
        default="",
        help="File to store the results. If not provided, the results will be stored in `output/results-<timestamp>.json`.",
    )
    parser.add_argument(
        "--time-format",
        default="gcomp",
        help="Time format to use, options are 'gcomp' and 'grift'",
        choices=["gcomp", "grift"],
    )
    args = parser.parse_args()
    config = Config(
        bench_directory=args.benchmarks_directory,
        output_directory=args.output_directory,
        compile_command=args.compile_cmd,
        run_command=args.run_cmd,
        file_ext=args.file_ext,
        num_runs=args.num_runs,
        use_internal_timing=args.internal_timing,
        output_file=args.output_file,
        time_format=args.time_format,
    )
    return config


def main():
    config = parse_args()
    print("Configuration:")
    print(f"  benchmarks_directory: {config.bench_directory}")
    print(f"  output_directory: {config.output_directory}")
    print(f"  compile_command: {config.compile_command}")
    print(f"  run_command: {config.run_command}")
    print(f"  num_runs: {config.num_runs}")
    print(f"  use_internal_timing: {config.use_internal_timing}")
    print(f"  output_file: {config.output_file}")
    print(f"  time_format: {config.time_format}")
    print(f"")
    results = compile_and_run_benchmarks(config)
    parsed_results = parse_results(config, results)
    write_results(config, parsed_results)


if __name__ == "__main__":
    main()
