import asyncio
import traceback

class Job:
    def __init__(self, data):
        self._id = None
        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):
        try:
            while True:
                job: Job = await self.input_queue.get()
                self._jobs_dequeued += 1
                if self.sequential_node == False:
                    async for job in self.process_job(job):
                        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
                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)
                        async for job in self.process_job(job):
                            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
                        self.next_i += 1
        except Exception as e:
            print(f"An error occurred in node: {self.__class__.__name__} worker: {self.worker_id}: {e}")
            traceback.print_exc()
            raise  # Re-raises the last exception.

    async def process_job(self, job: Job):
        raise NotImplementedError()
    
class Pipeline:
    def __init__(self):
        self.input_queues = []
        self.root_queue = None
        # self.output_queues = []
        # self.job_sysncs = []
        self.nodes= []
        self.node_workers = {}
        self.tasks = []
        self._job_id = 0

    async def add_node(self, node: Node, num_workers=1, input_queue=None, output_queue=None, job_sync=None, sequential_node=False ):
        # input_queue must not be None
        if input_queue is None:
            raise ValueError('input_queue is None')
        # job_sync nodes must be sequential_nodes
        if job_sync is not None and sequential_node == False:
            raise ValueError('job_sync is not None and sequential_node is False')
        # sequential_nodes should one have 1 worker
        if sequential_node == True and num_workers != 1:
            raise ValueError('sequentaial nodes can only have one node (sequential_node is True and num_workers is not 1)')
        # output queue must not equal input_queue
        if output_queue == input_queue:
            raise ValueError('output_queue must not be the same as input_queue')        

        node_name = node.__name__
        if node_name not in self.nodes:
            self.nodes.append(node_name)

        # if input_queue is None then this is the root node
        if len(self.input_queues) == 0:
            self.root_queue = input_queue

        self.input_queues.append(input_queue)

        for i in range(num_workers):
            worker_id = i
            node_worker = node(worker_id, input_queue, output_queue, job_sync, sequential_node)
            if node_name not in self.node_workers:
                self.node_workers[node_name] = []
            self.node_workers[node_name].append(node_worker)
            task = asyncio.create_task(node_worker.run())
            self.tasks.append(task)

    async def enqueue_job(self, job: Job):
        job.id = self._job_id
        self._job_id += 1
        await self.root_queue.put(job)

    async def close(self):
        for task in self.tasks:
            task.cancel()
        await asyncio.gather(*self.tasks, return_exceptions=True)