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")


async def child(scope: structio.TaskScope, x: float):
    print(f"[main] Child is alive! Canceling scope in {x} seconds")
    await structio.sleep(x)
    scope.cancel()


async def main_child(x: float):
    print("[main] Parent is alive")
    t = structio.clock()
    async with structio.create_pool() as p:
        print("[main] Spawning child")
        p.spawn(child, p.scope, x / 2)
        print(f"[main] Child spawned, sleeping for {x} seconds")
        await structio.sleep(x)
    print(f"[main] Done in {structio.clock() - t:.2f} seconds")

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