import structio


async def sleeper(n):
    print(f"[sleeper] Going to sleep for {n} seconds!")
    i = structio.clock()
    try:
        await structio.sleep(n)
    except structio.Cancelled:
        print(f"[sleeper] Oh no, I've been cancelled! (was gonna sleep {structio.clock() - i:.2f} more seconds)")
        raise
    print("[sleeper] Woke up!")


async def main_simple(n, o, p):
    print(f"[main] Parent is alive, spawning {o} children sleeping {n} seconds each")
    t = structio.clock()
    async with structio.create_pool() as pool:
        for i in range(o):
            pool.spawn(sleeper, n)
        print(f"[main] Children spawned, sleeping {p} seconds before cancelling")
        await structio.sleep(p)
        # Note that cancellations propagate to all inner task scopes!
        pool.scope.cancel()
    print(f"[main] Parent exited in {structio.clock() - t:.2f} seconds")


async def main_nested(n, o, p):
    print(f"[main] Parent is alive, spawning {o} children in two contexts sleeping {n} seconds each")
    t = structio.clock()
    async with structio.create_pool() as p1:
        for i in range(o):
            p1.spawn(sleeper, n)
        async with structio.create_pool() as p2:
            for i in range(o):
                p2.spawn(sleeper, n)
            print(f"[main] Children spawned, sleeping {p} seconds before cancelling")
            await structio.sleep(p)
            # Note that cancellations propagate to all inner task scopes!
            p1.scope.cancel()
    print(f"[main] Parent exited in {structio.clock() - t:.2f} seconds")


# Should take about 4 seconds
structio.run(main_simple, 5, 2, 2)
structio.run(main_nested, 5, 2, 2)