#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   graceful_exit.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK graceful exit utilities.
"""

import signal
import threading
from collections.abc import Generator
from contextlib import contextmanager
from typing import Any


class GracefulExit:
    """Handle graceful exit on keyboard interrupt.

    Supports chained handlers where child handlers automatically inherit
    cancellation state from their parent. This is useful for nested operations
    where cancellation should propagate down the call chain.

    Example:
        with graceful_exit("Upload cancelled") as parent:
            child = parent.create_child("Part upload cancelled")
            # If parent.exit_now becomes True, child.exit_now also returns True

    """

    def __init__(
        self,
        message: str = "Operation cancelled by user",
        parent: "GracefulExit | None" = None,
    ):
        """Initialize the graceful exit handler.

        Args:
            message: Message to display when interrupted.
            parent: Optional parent handler to chain cancellation from.

        """
        self.message = message
        self._exit_now = False
        self._parent = parent
        self.original_sigint = None

    @property
    def exit_now(self) -> bool:
        """Check if exit has been requested.

        Returns True if this handler or any parent in the chain has been cancelled.
        """
        if self._exit_now:
            return True
        if self._parent is not None:
            return self._parent.exit_now
        return False

    @exit_now.setter
    def exit_now(self, value: bool) -> None:
        """Set the exit flag."""
        self._exit_now = value

    def create_child(self, message: str | None = None) -> "GracefulExit":
        """Create a child handler chained to this parent.

        The child handler will inherit cancellation state from this parent.
        When the parent is cancelled, the child's exit_now will also return True.

        Args:
            message: Optional custom message for the child handler.
                If not provided, inherits the parent's message.

        Returns:
            A new GracefulExit instance chained to this handler.

        """
        return GracefulExit(
            message=message or self.message,
            parent=self,
        )

    def signal_handler(self, _signum: int, _frame: Any) -> None:
        """Handle SIGINT signal gracefully.

        Args:
            _signum: Signal number.
            _frame: Current stack frame.

        """
        self._exit_now = True
        # Restore original handler and re-raise to let KeyboardInterrupt propagate naturally
        if self.original_sigint is not None:
            signal.signal(signal.SIGINT, self.original_sigint)
        raise KeyboardInterrupt()

    def __enter__(self):
        """Enter the context manager.

        Returns:
            Self for use in with statements.

        """
        # Only set up signal handlers in the main thread
        if threading.current_thread() == threading.main_thread():
            self.original_sigint = signal.signal(signal.SIGINT, self.signal_handler)
        else:
            self.original_sigint = None
        return self

    def __exit__(self, _exc_type, _exc_val, _exc_tb):
        """Exit the context manager.

        Args:
            _exc_type: Exception type.
            _exc_val: Exception value.
            _exc_tb: Exception traceback.

        """
        # Restore original handler if not already restored by signal_handler
        # Only if we actually set up signal handlers (i.e., in main thread)
        if (
            self.original_sigint is not None
            and signal.getsignal(signal.SIGINT) == self.signal_handler
        ):
            signal.signal(signal.SIGINT, self.original_sigint)


@contextmanager
def graceful_exit(
    message: str = "Operation cancelled by user",
) -> Generator[GracefulExit, None, None]:
    """Context manager for graceful keyboard interrupt handling.

    Args:
        message: Message to display when interrupted

    Example:
        with graceful_exit("Download cancelled by user") as handler:
            for i in range(100):
                if handler.exit_now:
                    break
                # Do work...
                time.sleep(0.1)

    """
    handler = GracefulExit(message)
    with handler:
        yield handler
