"""Benchmark mypy on N distinct Literal overloads."""
import sys
import os
import tempfile
import time

def generate(n: int, workdir: str) -> tuple[str, str]:
    pkgdir = os.path.join(workdir, f"pkg_{n}")
    os.makedirs(pkgdir, exist_ok=True)
    rt = ["str", "int", "float"]
    lines = ["import typing", "@typing.overload"]
    for i in range(n):
        lines.append(f"def get_param(param: typing.Literal['overload_{i}_{rt[i%3]}']) -> {rt[i%3]}: ...")
        if i < n - 1:
            lines.append("@typing.overload")
    lines.append("def get_param(param: str) -> object: ...")
    with open(os.path.join(pkgdir, "__init__.pyi"), "w") as f:
        f.write("\n".join(lines))
    open(os.path.join(pkgdir, "py.typed"), "w").close()

    script = os.path.join(workdir, f"script_{n}.py")
    with open(script, "w") as f:
        f.write(f"import sys\n")
        f.write(f"sys.path.insert(0, {repr(pkgdir + '/..')})\n")
        f.write(f"from pkg_{n} import get_param\n")
        f.write(f"x = get_param('overload_0_str')\n")
    cache = os.path.join(workdir, f"cache_{n}")
    return script, cache

def bench(n: int, workdir: str) -> float:
    from mypy import api as mypy_api
    script, cache = generate(n, workdir)
    args = [
        script,
        "--no-error-summary",
        f"--cache-dir={cache}",
        "--no-incremental",
        "--python-version=3.11",
    ]
    start = time.perf_counter()
    mypy_api.run(args)
    return time.perf_counter() - start

if __name__ == "__main__":
    sizes = list(map(int, sys.argv[1:])) if len(sys.argv) > 1 else [500, 1000]
    with tempfile.TemporaryDirectory() as workdir:
        for n in sizes:
            t = bench(n, workdir)
            print(f"n={n:6d}: {t:.2f}s")
            sys.stdout.flush()