#!/usr/bin/env python3
"""Where the tree store's extra sync work goes (docs/research/084-parallel-syncs.md).

Builds pure-serve with three stores from store/store-tree.c: as it is; with no
compaction; and with no compaction and a 64 MiB tail written at creation (no
tail growth during a run). Then durable score posts (64 connections), the
builds in rotated order per round, each on a fresh store after 64 warm-up
posts. Prints one JSON line per run and a summary.
"""
import argparse
import json
import select
import shutil
import statistics
import subprocess
import sys
import tempfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2]
HTTP = ROOT / 'benchmarks' / 'http'
PATCHES = {
    'as it is': [],
    'no compaction': [('enum { COMPACT_SLACK = 1048576 };', 'enum { COMPACT_SLACK = 1 << 30 };')],
    'no compaction, no growth': [('enum { COMPACT_SLACK = 1048576 };', 'enum { COMPACT_SLACK = 1 << 30 };'),
                                 ('TAIL_STEP = 262144', 'TAIL_STEP = 67108864')],
}


def build(work, name, patches):
    source = (ROOT / 'store' / 'store-tree.c').read_text()
    for old, new in patches:
        assert old in source, old
        source = source.replace(old, new)
    store = work / f'store-{len(name)}.c'
    store.write_text(source)
    binary = work / f'pure-serve-{len(name)}'
    subprocess.run(['cc', '-O2', '-std=c17', f'-I{ROOT / "core"}', f'-I{ROOT / "serve"}', '-DPURE_ASYNC_EFFECTS', '-pthread',
                    *[str(ROOT / s) for s in ('serve/serve.c', 'serve/deliver.c', 'serve/message-json.c', 'core/module.c',
                                              'core/module-list.c', 'core/module-save.c', 'core/module-text.c',
                                              'core/module-store.c')], str(store), str(ROOT / 'core/runtime.c'),
                    '-lssl', '-lcrypto', '-lm', '-o', str(binary)], check=True)
    return binary


def posts(binary, module, directory, seconds):
    store = Path(tempfile.mkdtemp(dir=directory, prefix='pure-store-variant-'))
    shutil.rmtree(store)
    process = subprocess.Popen(['taskset', '-c', '0-2', str(binary), '--module', str(module), '--store', str(store),
                                '--port', '0'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        ready, _, _ = select.select([process.stdout], [], [], 20)
        assert ready, 'pure-serve did not start'
        port = json.loads(process.stdout.readline())['port']
        load = ['taskset', '-c', '3', str(HTTP / 'load'), str(port), '64']
        subprocess.run(load + ['1', 'score', '64', 'keep-alive', '16'], capture_output=True, check=True)
        measured = json.loads(subprocess.run(load + [str(seconds), 'score', '64', 'keep-alive', '16'], capture_output=True,
                                             text=True, check=True).stdout)
        return {'posts_per_second': measured['2xx'] / seconds, 'p50_us': measured['p50_us'], 'p99_us': measured['p99_us'],
                'max_us': measured['max_us'], 'errors': measured['errors'] + measured['4xx'] + measured['5xx']}
    finally:
        process.terminate()
        process.wait()
        shutil.rmtree(store, ignore_errors=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--directory', type=Path, default=Path('/var/tmp'), help='on the device to measure')
    parser.add_argument('--aup', type=Path, default=ROOT.parent / 'aup' / 'build' / 'aup')
    parser.add_argument('--rounds', type=int, default=6)
    parser.add_argument('--seconds', type=float, default=4)
    args = parser.parse_args()
    with tempfile.TemporaryDirectory() as work:
        work = Path(work)
        module = work / 'scores.pmod'
        subprocess.run([str(args.aup), 'build', str(HTTP / 'scores.aup'), '-o', str(module)], check=True)
        subprocess.run(['cc', '-O2', '-std=c17', '-o', str(HTTP / 'load'), str(HTTP / 'load.c')], check=True)
        binaries = {name: build(work, name, patches) for name, patches in PATCHES.items()}
        names, runs = list(binaries), {name: [] for name in binaries}
        for round_ in range(args.rounds):
            for name in names[round_ % len(names):] + names[:round_ % len(names)]:
                measured = posts(binaries[name], module, args.directory, args.seconds)
                runs[name].append(measured)
                print(json.dumps({'round': round_, 'store': name, **measured}), flush=True)
        for name in names:
            print(json.dumps({'store': name, 'rounds': args.rounds,
                              'posts_per_second_median': statistics.median(r['posts_per_second'] for r in runs[name]),
                              'p99_us_median': statistics.median(r['p99_us'] for r in runs[name]),
                              'max_us': max(r['max_us'] for r in runs[name])}))


if __name__ == '__main__':
    sys.exit(main())