from __future__ import annotations

import ast
import pyperf


def make_src(n, chunk_len, per_line: int = 64):
    assert n > 0 and chunk_len >= 0 and per_line > 0
    chunk = "b'" + ("x" * chunk_len) + "'"
    parts = [chunk] * n
    lines = ["x = ("]
    while parts:
        group = " ".join(parts[:per_line])
        parts = parts[per_line:]
        lines.append(f"    {group}")
    lines.append(")")
    return "\n".join(lines)


def bench_compile(loops, n, chunk_len, per_line: int = 64):
    src = make_src(n, chunk_len, per_line)
    t0 = pyperf.perf_counter()
    for _ in range(loops):
        compile(src, "<bench>", "exec")
    return pyperf.perf_counter() - t0


def bench_ast_parse(loops, n, chunk_len, per_line: int = 64):
    src = make_src(n, chunk_len, per_line)
    t0 = pyperf.perf_counter()
    for _ in range(loops):
        ast.parse(src, filename="<bench>", mode="exec")
    return pyperf.perf_counter() - t0


def main():
    runner = pyperf.Runner()

    for n in (1, 2, 4, 8, 16, 32, 64, 128, 1024):
        runner.bench_time_func(
            f"compile_bytes_concat_n{n}_chunk1",
            bench_compile,
            n,
            1,
        )
        runner.bench_time_func(
            f"parse_bytes_concat_n{n}_chunk1",
            bench_ast_parse,
            n,
            1,
        )

    for n, chunk in ((256, 4), (4, 128), (4, 256)):
        runner.bench_time_func(
            f"compile_bytes_concat_n{n}_chunk{chunk}",
            bench_compile,
            n,
            chunk,
        )
        runner.bench_time_func(
            f"parse_bytes_concat_n{n}_chunk{chunk}",
            bench_ast_parse,
            n,
            chunk,
        )


if __name__ == "__main__":
    main()