"""Count primes with an import-safe process pool. Python 3.9+, standard library."""
import argparse
from concurrent.futures import ProcessPoolExecutor
import json
from math import isqrt
import multiprocessing as mp
import sys


def count_chunk(bounds):
    start, stop = bounds
    if start < 2 or stop < start:
        raise ValueError("invalid interval")
    count = 0
    for number in range(start, stop):
        if number == 2 or (number % 2 and all(
            number % divisor for divisor in range(3, isqrt(number) + 1, 2)
        )):
            count += 1
    return count


def calculate(chunks, workers, serial=False):
    if serial:
        return sum(map(count_chunk, chunks))
    with ProcessPoolExecutor(
        max_workers=workers, mp_context=mp.get_context("spawn")
    ) as pool:
        return sum(pool.map(count_chunk, chunks))


def main():
    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
    parser.add_argument("--limit", type=int, default=10000)
    parser.add_argument("--chunk-size", type=int, default=2500)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--serial", action="store_true", help="skip the pool")
    args = parser.parse_args()
    if not 2 <= args.limit <= 1000000:
        parser.error("--limit must be 2..1000000")
    if not 1 <= args.chunk_size <= 1000000:
        parser.error("--chunk-size must be 1..1000000")
    if not 1 <= args.workers <= 8:
        parser.error("--workers must be 1..8")
    task_count = (args.limit - 1 + args.chunk_size - 1) // args.chunk_size
    if task_count > 1000:
        parser.error("too many chunks; increase --chunk-size (maximum 1000 chunks)")
    chunks = [
        (start, min(start + args.chunk_size, args.limit + 1))
        for start in range(2, args.limit + 1, args.chunk_size)
    ]
    try:
        total = calculate(chunks, args.workers, args.serial)
    except (OSError, RuntimeError, ValueError) as exc:
        print(f"ERROR: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1
    print(json.dumps({"limit": args.limit, "chunks": len(chunks),
                      "prime_count": total}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())