import concurrent.futures
import multiprocessing
import os
import pickle
import sys
import time

from pathlib import Path
from typing import Iterator


RUNNERS = 100
barrier = multiprocessing.Barrier(RUNNERS)

VERBOSE = False


def get_pickles() -> Iterator[Path]:
    """
    List all pickle files in the lib2to3 folder.
    """
    # Assume that the lib2to3 code is in the same directory as the multiprocessing package.
    standard_packages_path = Path(multiprocessing.__file__).parent.parent

    lib2to3_package_path = standard_packages_path / "lib2to3"
    assert lib2to3_package_path.is_dir()

    return lib2to3_package_path.glob("**/*.pickle")


def do_import(iteration: int) -> None:
    # Verify that the lib2to3 modules are not imported yet.
    assert "lib2to3.patcomp" not in sys.modules
    assert "lib2to3.pygram" not in sys.modules

    # First get all the processes to be ready.
    barrier.wait()
    # Then let them run with just a tiny bit of offset.
    time.sleep(iteration / 100)

    import lib2to3.refactor

    # Verify that the lib2to3 modules are actually imported now.
    assert "lib2to3.patcomp" in sys.modules
    assert "lib2to3.pygram" in sys.modules


def main():
    # Remove the existing pickle files.
    for p in get_pickles():
        os.unlink(p)

    all_futures = []
    with concurrent.futures.ProcessPoolExecutor(max_workers=RUNNERS) as executor:
        for i in range(RUNNERS):
            all_futures.append(executor.submit(do_import, i))

    failed = 0
    for i, future in enumerate(all_futures):
        try:
            future.result()
        except (EOFError, pickle.UnpicklingError):
            failed += 1

    assert list(get_pickles()), (
        "No pickle files found after using lib2to3! "
        "Did you run this with a non-writable Python installation?"
    )

    if failed:
        print("FAILED: {} out of {}".format(failed, RUNNERS))
    else:
        print("No failure")


if __name__ == "__main__":
    main()