import structio
import time


def fake_async_sleeper(n, name: str = ""):
    print(f"[thread{f' {name}' if name else ''}] About to sleep for {n} seconds")
    t = time.time()
    if structio.thread.is_async_thread():
        print(f"[thread{f' {name}' if name else ''}] I have async superpowers!")
        structio.thread.run_coro(structio.sleep, n)
    else:
        print(f"[thread{f' {name}' if name else ''}] Using old boring time.sleep :(")
        time.sleep(n)
    print(
        f"[thread{f' {name}' if name else ''}] Slept for {time.time() - t:.2f} seconds"
    )
    return n


async def main(n):
    print(f"[main] Spawning worker thread, exiting in {n} seconds")
    t = structio.clock()
    d = await structio.thread.run_in_worker(fake_async_sleeper, n)
    assert d == n
    print(f"[main] Exited in {structio.clock() - t:.2f} seconds")


async def main_timeout(n, k):
    print(f"[main] Spawning worker thread, exiting in {k} seconds")
    t = structio.clock()
    with structio.skip_after(k):
        # We need to make the operation explicitly cancellable if we want
        # to be able to move on!
        await structio.thread.run_in_worker(fake_async_sleeper, n, cancellable=True)
    print(f"[main] Exited in {structio.clock() - t:.2f} seconds")


async def main_multiple(n, k):
    print(f"[main] Spawning {n} worker threads each sleeping for {k} seconds")
    t = structio.clock()
    async with structio.create_pool() as pool:
        for i in range(n):
            pool.spawn(structio.thread.run_in_worker, fake_async_sleeper, k, str(i))
        print(f"[main] Workers spawned")
    # Keep in mind that there is some overhead associated with running worker threads,
    # not to mention that it gets tricky with how the OS schedules them and whatnot. So,
    # it's unlikely that all threads finish exactly at the same time and that we exit in
    # k seconds, even just because there's a lot of back and forth going on under the hood
    # between structio and the worker threads themselves
    print(f"[main] Exited in {structio.clock() - t:.2f} seconds")


structio.run(main, 2)
structio.run(main_timeout, 5, 3)
structio.run(main_multiple, 10, 2)