import asyncio
import threading
import queue
import time
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Callable, Coroutine

class ComfyUIAsyncHandler:
    """Handles async operations in ComfyUI's synchronous execution context"""

    def __init__(self, max_workers: int = 4):
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.active_tasks = {}
        self.logger = logging.getLogger(__name__)

    def run_async_in_thread(self, coro: Coroutine) -> Any:
        """Run async coroutine in separate thread with event loop"""

        def run_in_thread():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                return loop.run_until_complete(coro)
            finally:
                loop.close()

        future = self.executor.submit(run_in_thread)
        return future.result()

    def run_with_progress(self, coro: Coroutine, progress_callback: Callable = None) -> Any:
        """Run async operation with progress tracking"""

        progress_queue = queue.Queue()
        result_queue = queue.Queue()

        def async_wrapper():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

            async def tracked_operation():
                task = asyncio.create_task(coro)

                # Monitor progress
                start_time = time.time()
                while not task.done():
                    if progress_callback:
                        elapsed = time.time() - start_time
                        progress_queue.put(f"Processing... {elapsed:.1f}s")
                    await asyncio.sleep(0.5)

                return await task

            try:
                result = loop.run_until_complete(tracked_operation())
                result_queue.put(("success", result))
            except Exception as e:
                self.logger.error(f"Async operation failed: {e}")
                result_queue.put(("error", e))
            finally:
                loop.close()

        # Start async operation in thread
        thread = threading.Thread(target=async_wrapper)
        thread.start()

        # Handle progress updates
        while thread.is_alive():
            try:
                progress = progress_queue.get_nowait()
                if progress_callback:
                    progress_callback(progress)
            except queue.Empty:
                time.sleep(0.1)

        thread.join()

        # Get result
        status, result = result_queue.get()
        if status == "error":
            raise result
        return result

    def shutdown(self):
        """Shutdown the executor"""
        self.executor.shutdown(wait=True)

# Global async handler instance
async_handler = ComfyUIAsyncHandler()