import structio


async def child(i: int, sem: structio.Semaphore):
    async with sem:
        print(f"[child {i}] Entered critical section")
        await structio.sleep(1)
    print(f"[child {i}] Exited critical section")


async def main_sem(n: int, k: int):
    assert isinstance(n, int) and n > 0
    assert isinstance(k, int) and k > 1
    print(f"[main] Parent is alive, creating semaphore of size {n}")
    semaphore = structio.Semaphore(n)
    t = structio.clock()
    async with structio.create_pool() as pool:
        print(f"[main] Spawning {n * k} children")
        for i in range(1, (n * k) + 1):
            pool.spawn(child, i, semaphore)
        print("[main] All children spawned, waiting for completion")
        # Since our semaphore has a limit of n tasks that
        # can acquire it concurrently, we should see at most
        # n instances of child running at any given time,
        # like in batches
    print(f"[main] Done in {structio.clock() - t:.2f} seconds")


async def main_lock(k: int):
    assert isinstance(k, int) and k > 1
    print(f"[main] Parent is alive, creating lock")
    lock = structio.Lock()
    t = structio.clock()
    async with structio.create_pool() as pool:
        print(f"[main] Spawning {k} children")
        for i in range(1, k + 1):
            # Locks are implemented as simple binary semaphores
            # and have an identical API, so they can be used
            # interchangeably
            pool.spawn(child, i, lock)
        print("[main] All children spawned, waiting for completion")
    print(f"[main] Done in {structio.clock() - t:.2f} seconds")


async def recursive_child(i: int, sem: structio.Semaphore, kapow: bool = False):
    async with sem:
        print(f"[{'copy of ' if kapow else ''}child {i}] Entered critical section")
        if kapow:
            await recursive_child(i, sem)
        await structio.sleep(1)
    print(f"[{'copy of ' if kapow else ''}child {i}] Exited critical section")


async def main_rlock(k: int):
    assert isinstance(k, int) and k > 1
    print(f"[main] Parent is alive, creating recursive lock")
    lock = structio.RLock()
    t = structio.clock()
    async with structio.create_pool() as pool:
        print(f"[main] Spawning {k} children")
        for i in range(1, k + 1):
            pool.spawn(recursive_child, i, lock, True)
        print("[main] All children spawned, waiting for completion")
    print(f"[main] Done in {structio.clock() - t:.2f} seconds")


# Should exit in about k seconds
structio.run(main_sem, 3, 5)
# Same here, should exit in k seconds, but
# it'll run one task at a time (also fewer tasks)
structio.run(main_lock, 10)
# This should exit in about 2k seconds
structio.run(main_rlock, 5)