import json
import sys

import pyperf
from memray_helper import get_tracker


EMPTY = ({}, 2000)
SIMPLE_DATA = {
    "key1": 0,
    "key2": True,
    "key3": "value",
    "key4": "foo",
    "key5": "string",
}
SIMPLE = (SIMPLE_DATA, 1000)
NESTED_DATA = {
    "key1": 0,
    "key2": SIMPLE[0],
    "key3": "value",
    "key4": SIMPLE[0],
    "key5": SIMPLE[0],
    "key": "\u0105\u0107\u017c",
}
NESTED = (NESTED_DATA, 1000)
HUGE = ([NESTED[0]] * 1000, 1)

CASES = ["EMPTY", "SIMPLE", "NESTED", "HUGE"]


def bench_json_dumps(data):
    with get_tracker():
        for obj, count_it in data:
            for _ in count_it:
                json.dumps(obj)


def add_cmdline_args(cmd, args):
    if args.cases:
        cmd.extend(("--cases", args.cases))


def main():
    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
    runner.argparser.add_argument(
        "--cases",
        help="Comma separated list of cases. Available cases: %s. By default, run all cases."
        % ", ".join(CASES),
    )
    runner.metadata["description"] = "Benchmark json.dumps()"

    args = runner.parse_args()
    if args.cases:
        cases = []
        for case in args.cases.split(","):
            case = case.strip()
            if case:
                cases.append(case)
        if not cases:
            print("ERROR: empty list of cases")
            sys.exit(1)
    else:
        cases = CASES

    data = []
    for case in cases:
        obj, count = globals()[case]
        data.append((obj, range(count)))

    runner.bench_func("json_dumps", bench_json_dumps, data)


if __name__ == "__main__":
    main()
