#!/usr/bin/env python3

import sys
import os
import subprocess
import re
from collections import namedtuple
import argparse

factorio_bin = "/home/redacted/factorio/versions/0.17.6/bin/x64/factorio"
factorio_config = (sys.path[0] or os.getcwd()) + "/benchconf.ini"
factorio_moddir = (sys.path[0] or os.getcwd()) + "/bench-moddir"

headless_bin = "/home/redacted/factorio/versions/0.17.6-headless/bin/x64/factorio"

def run_benchmark(save_file_name, ticks, headless=False):
    """
    Get a benchmark Result(avg, min, max) ms per update.
    """
    if headless:
        command = [ headless_bin,
                    "--mod-directory", factorio_moddir,
                    "--benchmark", save_file_name,
                    "--benchmark-ticks", str(ticks)]
    else:
        command = [ factorio_bin,
                    "--disable-audio",
                    "--config", factorio_config,
                    "--mod-directory", factorio_moddir,
                    "--benchmark", save_file_name,
                    "--benchmark-ticks", str(ticks) ]
    try:
        result_str = subprocess.check_output(command).decode()
    except subprocess.CalledProcessError as e:
        print(e.output.decode())
        exit(e.returncode)
    #total time
    match = re.search('Performed \d+ updates in (\d+\.\d+) ms', result_str)
    total = float(match.group(1))
    #statistics
    match = re.search(
            "avg: (\d+\.\d+) ms, min: (\d+\.\d+) ms, max: (\d+\.\d+) ms",
            result_str)
    avg, low, high = (float(s) for s in match.groups())
    #result
    return namedtuple("Result", ["total", "avg", "min", "max"])(
            total, avg, low, high)

def best_of_n(save_file_name, ticks, n_runs, headless=False):
    results = [run_benchmark(save_file_name, ticks, headless)
            for i in range(n_runs)]
    results.sort()
    return results[0]

def pretty_format_result(result):
    fraction_realtime = 1000/60/result.avg
    return "{:7.3f} × realtime, avg={:.3f} min={:.3f} max={:.3f}".format(
            fraction_realtime, result.avg, result.min, result.max)

def main():
    parser = argparse.ArgumentParser(
            description="Benchmark factorio save files",
            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("file", type=str, help="save file name")
    parser.add_argument("--ticks",
            type=int, metavar="M", default=1800, help="length of benchmark")
    parser.add_argument("--best-of",
            type=int, metavar="N", default=3, help="best of N runs")
    parser.add_argument("--headless",
            action="store_true", help="use the headless binary")
    args = parser.parse_args()
    result = best_of_n(args.file, args.ticks, args.best_of, args.headless)
    print(args.file + ": " + pretty_format_result(result))

if __name__ == "__main__":
    main()