import time
import structio


async def producer(q: structio.Queue, n: int):
    for i in range(n):
        # This will wait until the
        # queue is emptied by the
        # consumer
        await q.put(i)
        print(f"[producer] Produced {i}")
    await q.put(None)
    print("[producer] Producer done")


async def consumer(q: structio.Queue):
    while True:
        # Hangs until there is
        # something on the queue
        item = await q.get()
        if item is None:
            print("[consumer] Consumer done")
            break
        print(f"[consumer] Consumed {item}")
        # Simulates some work so the
        # producer waits before putting
        # the next value
        await structio.sleep(1)


def threaded_consumer(q: structio.thread.AsyncThreadQueue):
    while True:
        # Hangs until there is
        # something on the queue
        item = q.get_sync()
        if item is None:
            print("[worker consumer] Consumer done")
            break
        print(f"[worker consumer] Consumed {item}")
        # Simulates some work so the
        # producer waits before putting
        # the next value
        time.sleep(1)
    return 69


async def main(q: structio.Queue, n: int):
    print("[main] Starting consumer and producer")
    t = structio.clock()
    async with structio.create_pool() as ctx:
        ctx.spawn(producer, q, n)
        ctx.spawn(consumer, q)
    print(f"[main] Exited in {structio.clock() - t:.2f} seconds")


def threaded_producer(q: structio.thread.AsyncThreadQueue, n: int):
    print("[worker producer] Producer started")
    for i in range(n):
        # This will wait until the
        # queue is emptied by the
        # consumer
        q.put_sync(i)
        print(f"[worker producer] Produced {i}")
    q.put_sync(None)
    print("[worker producer] Producer done")
    return 42


async def main_threaded(q: structio.thread.AsyncThreadQueue, n: int):
    print("[main] Starting consumer and producer")
    t = structio.clock()
    async with structio.create_pool() as pool:
        pool.spawn(producer, q, n)
        val = await structio.thread.run_in_worker(threaded_consumer, q)
        print(f"[main] Thread returned {val!r}")
    print(f"[main] Exited in {structio.clock() - t:.2f} seconds")


async def main_threaded_2(q: structio.thread.AsyncThreadQueue, n: int):
    print("[main] Starting consumer and producer")
    t = structio.clock()
    async with structio.create_pool() as pool:
        pool.spawn(consumer, q)
        val = await structio.thread.run_in_worker(threaded_producer, q, n)
        print(f"[main] Thread returned {val!r}")
    print(f"[main] Exited in {structio.clock() - t:.2f} seconds")


if __name__ == "__main__":
    queue = structio.Queue(2)  # Queue has size limit of 2
    structio.run(main, queue, 5)
    queue = structio.thread.AsyncThreadQueue(2)
    structio.run(main_threaded, queue, 5)
    structio.run(main_threaded_2, queue, 5)