import structio


async def successful(name: str, n):
    before = structio.clock()
    print(f"[child {name}] Sleeping for {n} seconds")
    await structio.sleep(n)
    print(f"[child {name}] Done! Slept for {structio.clock() - before:.2f} seconds")
    return n


async def failing(name: str, n):
    before = structio.clock()
    print(f"[child {name}] Sleeping for {n} seconds")
    await structio.sleep(n)
    print(f"[child {name}] Done! Slept for {structio.clock() - before:.2f} seconds, raising now!")
    raise TypeError("waa")


async def main(
    children_outer: list[tuple[str, int]], children_inner: list[tuple[str, int]]
):
    before = structio.clock()
    try:
        async with structio.create_pool() as p1:
            print(f"[main] Spawning children in first context ({hex(id(p1))})")
            for name, delay in children_outer:
                p1.spawn(successful, name, delay)
            print("[main] Children spawned")
            async with structio.create_pool() as p2:
                print(f"[main] Spawning children in second context ({hex(id(p2))})")
                for name, delay in children_inner:
                    p2.spawn(failing, name, delay)
                print("[main] Children spawned")
    except TypeError:
        print("[main] TypeError caught!")
    print(f"[main] Children exited in {structio.clock() - before:.2f} seconds")


async def main_nested(
    children_outer: list[tuple[str, int]], children_inner: list[tuple[str, int]]
):
    before = structio.clock()
    try:
        async with structio.create_pool() as p1:
            print(f"[main] Spawning children in first context ({hex(id(p1))})")
            for name, delay in children_outer:
                p1.spawn(successful, name, delay)
            print("[main] Children spawned")
            async with structio.create_pool() as p2:
                print(f"[main] Spawning children in second context ({hex(id(p2))})")
                for name, delay in children_outer:
                    p2.spawn(successful, name, delay)
                print("[main] Children spawned")
                async with structio.create_pool() as p3:
                    print(f"[main] Spawning children in third context ({hex(id(p3))})")
                    for name, delay in children_inner:
                        p3.spawn(failing, name, delay)
                    print("[main] Children spawned")
    except TypeError:
        print("[main] TypeError caught!")
    print(f"[main] Children exited in {structio.clock() - before:.2f} seconds")

if __name__ == "__main__":
    structio.run(
        main,
        [("first", 1), ("third", 3)],
        [("second", 2), ("fourth", 4)],
    )
    structio.run(
        main_nested,
        [("first", 1), ("third", 3)],
        [("second", 2), ("fourth", 4)],
    )