import asyncio
import logging
from typing import Any, Callable

logger = logging.getLogger(__name__)


class AsyncTaskQueue:
    def __init__(self, worker_count: int = 3):
        self.queue = asyncio.Queue()
        self.worker_count = worker_count
        self.workers = []
        self._stop_event = asyncio.Event()

    async def start_workers(self):
        """Start the worker pool."""
        logger.debug(f"Starting {self.worker_count} workers...")
        for i in range(self.worker_count):
            worker = asyncio.create_task(self.worker_loop(f"Worker-{i+1}"))
            self.workers.append(worker)
        logger.debug("All workers started.")

    async def stop_workers(self):
        """Stop the worker pool and wait for all tasks to complete."""
        logger.debug("Stopping workers...")
        self._stop_event.set()  # Signal workers to stop
        await self.queue.join()  # Wait for all tasks to be processed
        for worker in self.workers:
            worker.cancel()
        await asyncio.gather(*self.workers, return_exceptions=True)
        logger.debug("All workers stopped.")


    async def worker_loop(self, name: str):
        """Worker loop: Process tasks from the queue."""
        logger.debug(f"{name} started.")
        while not self._stop_event.is_set():
            task = None
            try:
                task = await self.queue.get()
                logger.info(f"{name} processing: {task['name']}")
                await task['func'](*task['args'], **task['kwargs'])            
            except asyncio.CancelledError:
                logger.debug(f"{name} received cancellation.")
                break
            except Exception as e:
                logger.error(f"Error in {name}: {e}", exc_info=True)
            finally:
                if task is not None:
                    self.queue.task_done()
        logger.debug(f"{name} exited.")



    async def enqueue(self, name: str, func: Callable, *args: Any, **kwargs: Any):
        """Add a task to the queue."""
        logger.debug(f"Enqueuing task: {name}")
        await self.queue.put({'name': name, 'func': func, 'args': args, 'kwargs': kwargs})



async def run_sync_task_in_thread(func: Callable, *args: Any, **kwargs: Any):
    """Run a synchronous function in a separate thread.

    Args:
        func: The synchronous function to run.
        *args: Arguments to pass to the function.
    """
    await asyncio.to_thread(func, *args, **kwargs)
    
    """
from asyncTaskQueue import AsyncTaskQueue, run_sync_task_in_thread
import asyncio

# Example async task
async def your_async_function(name: str, duration: int):
    print(f"{name} started, will sleep for {duration} seconds")
    await asyncio.sleep(duration)
    print(f"{name} finished")

# Example sync task
def your_sync_function(name: str, duration: int):
    print(f"{name} started, will sleep for {duration} seconds")
    import time
    time.sleep(duration)
    print(f"{name} finished")

async def main():
    queue = AsyncTaskQueue(worker_count=2)
    await queue.start_workers()

    # Enqueue async tasks directly
    await queue.enqueue("AsyncTask1", your_async_function, "AsyncTask1", 2)
    await queue.enqueue("AsyncTask2", your_async_function, "AsyncTask2", 1)

    # Enqueue sync tasks using the wrapper
    await queue.enqueue("SyncTask1", run_sync_task_in_thread, your_sync_function, "SyncTask1", 2)

    # Let tasks run for a while
    await asyncio.sleep(5)

    # Stop workers
    await queue.stop_workers()

asyncio.run(main())


"""