import asyncio
import functools
import logging
import signal
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack, contextmanager
from threading import Event, Lock
from time import sleep
from typing import Any, Callable

logger = logging.getLogger()


def handler(signum, frame):
    print('Signal handler called with signal', signum)
    raise OSError("Couldn't open device!")


class AppDaemonRunContext:
    _stack: ExitStack
    loop: asyncio.AbstractEventLoop
    executor: ThreadPoolExecutor
    stop_event: Event
    shutdown_lock: Lock

    def __init__(self):
        self._stack = ExitStack()
        self.stop_event = Event()
        self.shutdown_lock = Lock()

    def __enter__(self):
        self.loop = self._stack.enter_context(self.asyncio_context())
        logger.debug("Entered asyncio context")

        signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
        for s in signals:
            self.loop.add_signal_handler(s, lambda s=s: self.loop.create_task(self.shutdown(s)))

        # self.loop.add_signal_handler(signal.SIGINT, self.handle_signal, signal.SIGINT)
        # self.loop.add_signal_handler(signal.SIGTERM, self.handle_signal)
        # self.executor = self._stack.enter_context(ThreadPoolExecutor(max_workers=5))
        self.executor = self._stack.enter_context(self.thread_context())
        logger.debug("Entered threadpool context")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._stack.__exit__(exc_type, exc_value, traceback)
        logger.debug("Exited context")

    async def shutdown(self, signal):
        """Cleanup tasks tied to the service's shutdown.
        
        https://www.roguelynn.com/words/asyncio-graceful-shutdowns/
        """
        logger.info(f"Received exit signal {signal.name}...")
        if not self.stop_event.is_set():
            logger.debug('Setting stop event')
            self.stop_event.set()
        
            tasks = [
                t for t in asyncio.all_tasks()
                if t is not asyncio.current_task()
            ]

            for task in tasks:
                # logger.debug(f"Waiting on task to finish: {task.get_coro().__qualname__}")
                logger.debug(f"Cancelling {task.get_name()}: {task.get_coro().__qualname__}")
                task.cancel()

            logger.debug('Waiting for tasks to finish...')
            try:
                await asyncio.gather(*tasks, return_exceptions=True)
            except asyncio.CancelledError as e:
                logger.debug(f'Cancelled: {e}')
            logger.debug("Stopping event loop in context shutdown")
            self.loop.stop()
        else:
            logger.warning('Already started shutdown')

    @contextmanager
    def asyncio_context(self):
        try:
            loop = asyncio.get_event_loop()
            yield loop
        finally:
            loop.close()
            if loop.is_closed():
                logger.debug("Closed the event loop.")

    @contextmanager
    def thread_context(self):
        with ThreadPoolExecutor(max_workers=5) as executor:
            yield executor
        if executor._shutdown:
            logger.debug('Shut down the ThreadPoolExecutor')

    async def run_in_executor(
        self,
        func: Callable,
        *args,
        timeout: float | None = None,
        **kwargs
    ) -> Any:
        """Run the sync function using the ThreadPoolExecutor and await the result"""
        timeout = timeout or 10.0

        coro = self.loop.run_in_executor(
            self.executor,
            functools.partial(func, **kwargs),
            *args,
        )

        try:
            return await asyncio.wait_for(coro, timeout)
        except asyncio.TimeoutError:
            print('Timed out')


if __name__ == "__main__":
    import logging
    import random
    from uuid import uuid4

    logging.basicConfig(level="DEBUG", format="{levelname:<8} {message}", style="{")

    def dummy_function(delay: float):
        id_ = uuid4().hex[:4]
        logger.info(f'{id_} sleeping for {delay:.1f}s')
        sleep(delay)
        logger.info(f'{id_} Done')

    async def async_dummy_function(delay: float):
        id_ = uuid4().hex[:4]
        logger.info(f'{id_} async sleeping for {delay:.1f}s')
        await asyncio.sleep(delay)
        logger.info(f'{id_} Done async')

    with AppDaemonRunContext() as cm:
        for _ in range(3):
            logger.info('Submitting random dummy_functions')
            cm.executor.submit(dummy_function, random.random() * 10.0)
            cm.loop.create_task(async_dummy_function(random.random() * 5.0))
        
        try:
            logger.info('Running until complete')
            cm.loop.run_until_complete(asyncio.gather(*asyncio.all_tasks(cm.loop)))
        except asyncio.CancelledError:
            logger.error('Cancelled')
    
    if cm.loop.is_closed():
        logger.debug('Loop is closed')
    if cm.executor._shutdown:
        logger.debug('Executor is shut down')