#!/usr/bin/env python3

import os
import re
import subprocess
import sys

def run_bench(cmd):
    """Run the benchmark with the given command (passed as an array of
    arguments).  The benchmark is expected to output on stdout one or more
    csv tables, of the form:

    === Table 5 ===

    protocol,request size,response size
    Open Invitation,249,543
    Trust Promotion,1677,2285
    Trust Migration,845,241
    Level Up,2253,241
    Issue Invite,1101,401
    Redeem Invite,1037,241
    Check Blockage,685,325
    Blockage Migration,973,241
    Update Invite,557,209
    Update Credential,749,209
    
    This function will return a dictionary with keys being the name of the
    table ("Table 5"), and the value being a list of csv rows, where each row
    is a dict with keys being the name of the column, and the value being the
    entry in that column."""
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, text=True)
    table_name = None
    seen_nonempty_line = False
    retval = {}
    table = []
    headers = []
    for line in proc.stdout:
        print(line,end='')
        line = line.rstrip()
        if matched := re.match(r'=== (.*) ===', line):
            if table_name is not None:
                # Store the previous table
                retval[table_name] = table
            table_name = matched.group(1)
            seen_nonempty_line = False
            table = []
            headers = []
            continue
        if line == "" and seen_nonempty_line == True:
            if table_name is not None:
                # The table is finished; store it
                retval[table_name] = table
                table_name = None
                seen_nonempty_line = False
                table = []
                headers = []
            continue
        if line == "":
            continue
        # At this point, we know the line is nonempty
        if len(headers) == 0:
            # This should be the header line
            headers = line.split(',')
            seen_nonempty_line = True
        else:
            values = line.split(',')
            row = dict(zip(headers, values))
            table.append(row)

    # If stdout ends, and there's a table in progress, store it
    if table_name is not None:
        retval[table_name] = table

    return retval

os.chdir(os.path.dirname(os.path.realpath(sys.argv[0])))
os.environ['PYTHONUNBUFFERED'] = '1'
if len(sys.argv) > 1:
    niters = sys.argv[1]
else:
    niters = "10"

lox_sigma_native = run_bench(["./run_native_bench", niters])
lox_zkp_native = run_bench(["./run_native_bench", "-z", niters])
lox_sigma_wasm = run_bench(["./run_wasm_bench", niters])
lox_zkp_wasm = run_bench(["./run_wasm_bench", "-z", niters])

print("\n*** Complete Table 2 ***\n")
print("protocol,zkp client native ms,zkp client wasm ms,zkp server native ms,sigma-rs client native ms,sigma-rs client wasm ms,sigma-rs server native ms")
for r in range(len(lox_sigma_native['Table 2'])):
    proto = lox_sigma_native['Table 2'][r]['protocol']
    assert proto == lox_zkp_native['Table 2'][r]['protocol']
    assert proto == lox_sigma_wasm['Table 2'][r]['protocol']
    assert proto == lox_zkp_wasm['Table 2'][r]['protocol']
    print(','.join([
        proto,
        lox_zkp_native['Table 2'][r]['client native ms'],
        lox_zkp_wasm['Table 2'][r]['client wasm ms'],
        lox_zkp_native['Table 2'][r]['server native ms'],
        lox_sigma_native['Table 2'][r]['client native ms'],
        lox_sigma_wasm['Table 2'][r]['client wasm ms'],
        lox_sigma_native['Table 2'][r]['server native ms'],
    ]))

print("\n*** Complete Table 5 ***\n")
print("protocol,zkp request bytes,zkp response bytes,sigma-rs request bytes,sigma-rs response bytes")
for r in range(len(lox_sigma_native['Table 5'])):
    proto = lox_sigma_native['Table 5'][r]['protocol']
    assert proto == lox_zkp_native['Table 5'][r]['protocol']
    print(','.join([
        proto,
        lox_zkp_native['Table 5'][r]['request size'],
        lox_zkp_native['Table 5'][r]['response size'],
        lox_sigma_native['Table 5'][r]['request size'],
        lox_sigma_native['Table 5'][r]['response size'],
    ]))