import asyncio
import random
import time


class Job:
    def __init__(self, id, data):
        self.id = id
        self.data = data


class Node:
    # def __init__(self, worker_id: int, input_queue, output_queue, buffer=None, job_sync=None):
    def __init__(self, worker_id: int, input_queue, output_queue=None, job_sync=None, sequential_node=False ):
        self.worker_id = worker_id
        self.input_queue = input_queue
        self.output_queue = output_queue
        self.buffer = {}
        self.job_sync = job_sync
        self.sequential_node = sequential_node
        self.next_i = 0
        self._jobs_dequeued = 0
        self._jobs_processed = 0
        # throw an error if job_sync is not None and sequential_node is False
        if self.job_sync is not None and self.sequential_node == False:
            raise ValueError('job_sync is not None and sequential_node is False')

    async def run(self):
        while True:
            job: Job = await self.input_queue.get()
            self._jobs_dequeued += 1
            if self.sequential_node == False:
                await self.process_job(job)
            else:
                # ensure that jobs are processed in order
                self.buffer[job.id] = job
                while self.next_i in self.buffer:
                    job = self.buffer.pop(self.next_i)
                    await self.process_job(job)
                    self.next_i += 1
            if self.output_queue is not None:
                await self.output_queue.put(job)
            if self.job_sync is not None:
                self.job_sync.append(job)
            self._jobs_processed += 1

    async def process_job(self, job: Job):
        raise NotImplementedError


class Node1(Node):
    async def process_job(self, job: Job):
        job.data += f' (processed by node 1, worker {self.worker_id})'


class Node2(Node):
    async def process_job(self, job: Job):
        sleep_duration = 0.8 + 0.4 * random.random()
        await asyncio.sleep(sleep_duration)
        job.data += f' (processed by node 2, worker {self.worker_id})'


class Node3(Node):
    async def process_job(self, job: Job):
        job.data += f' (processed by node 3, worker {self.worker_id})'
        print(f'{job.id} - {job.data}')


async def main():
    node1_queue = asyncio.Queue()
    node2_queue = asyncio.Queue()
    node3_queue = asyncio.Queue()

    num_jobs = 100
    joe_source = [Job(i, "") for i in range(num_jobs)]
    job_sync = []

    # create the workers
    num_workers = 5
    node1_workers = [Node1(i + 1, node1_queue, node2_queue) for i in range(1)]
    node2_workers = [Node2(i + 1, node2_queue, node3_queue) for i in range(num_workers)]
    node3_workers = [Node3(i + 1, node3_queue, job_sync=job_sync, sequential_node=True) for i in range(1)]

    # create tasks for the workers
    tasks1 = [asyncio.create_task(worker.run()) for worker in node1_workers]
    tasks2 = [asyncio.create_task(worker.run()) for worker in node2_workers]
    tasks3 = [asyncio.create_task(worker.run()) for worker in node3_workers]

    for job in joe_source:
        await node1_queue.put(job)
    # await input_queue.put(joe_source[0])

    try:
        while len(job_sync) < num_jobs:
            # print(f"Waiting for jobs to finish... Job sync size: {len(job_sync)}, node1_queue size: {node1_queue.qsize()}, node2_queue size: {node2_queue.qsize()}, node3_queue size: {node3_queue.qsize()}")
            await asyncio.sleep(0.1)
    except asyncio.CancelledError:
        print("Pipeline cancelled")
    for task in tasks1:
        task.cancel()
    for task in tasks2:
        task.cancel()
    for task in tasks3:
        task.cancel()
    await asyncio.gather(*tasks1, *tasks2, *tasks3, return_exceptions=True)


start_time = time.time()

try:
    asyncio.run(main())
except KeyboardInterrupt:
    print("Pipeline interrupted by user")

end_time = time.time()
print(f"Pipeline processed in {end_time - start_time} seconds.")