import asyncio
import signal
import sys
from common.logger import xlogger
from types import FrameType

SHUTTING_DOWN: bool = False


def signal_handler(*_):
    """Signal handler for main function. Run before uvicorn starts."""

    global SHUTTING_DOWN

    if SHUTTING_DOWN:
        return

    xlogger.warning("Shutdown signal called. Exiting gracefully.")
    SHUTTING_DOWN = True

    # Run async unloads for model
    # If an event loop doesn't exist (synchronous), exit.
    try:
        loop = asyncio.get_running_loop()
        loop.run_until_complete(signal_handler_async())
    except RuntimeError:
        sys.exit(0)


async def signal_handler_async():
    from common import model

    """
    Internal signal handler. Runs all async code to shut down the program.

    asyncio.run will cancel all remaining tasks and close the event loop.
    """

    if model.container:
        await model.unload_model(skip_wait=True, shutdown=True)

    if model.embeddings_container:
        await model.unload_embedding_model()


def uvicorn_signal_handler(signal_event: signal.Signals):
    """Overrides uvicorn's signal handler."""

    default_signal_handler = signal.getsignal(signal_event)

    def wrapped_handler(signum: int, frame: FrameType = None):
        xlogger.warning("Shutdown signal called. Exiting gracefully.")
        default_signal_handler(signum, frame)

    signal.signal(signal_event, wrapped_handler)