import traceback
from sys import exit
from typing import Type, Callable, Any
from contextlib import contextmanager


@contextmanager
def safe_run(debug: bool = False, exit_on_error: bool = True):
    """
    Context manager for safe code execution, catching exceptions and managing program termination.

    It specifically intercepts `KeyboardInterrupt` to exit with status 0, ensuring graceful shutdown
    when the user interrupts the script (e.g., Ctrl+C).

    :param debug: If True, prints the full traceback upon error. If False, prints a brief error message.
                  Defaults to False.
    :param exit_on_error: If True, the program exits with a status code of 1 upon catching an exception.
                          If False, the program continues execution after the 'with' block.
                          Defaults to True.
    :raises Exception: Catches all general exceptions (except KeyboardInterrupt, which is handled).
    :return: None.
    """
    try:
        yield
    except KeyboardInterrupt:
        exit(0)
    except Exception as e:
        if debug:
            traceback.print_exc()
        else:
            print(f'Error: {e}')
        if exit_on_error:
            exit(1)


def try_until_ok(
        func: Callable[..., Any],
        *args: Any,
        exceptions: tuple[Type[BaseException], ...] | Type[BaseException] = Exception,
        on_exception: str | Callable[[BaseException], Any] | None = None,
        **kwargs
        ) -> Any:
    """
    Attempts to execute the function repeatedly in a loop until it completes without errors,
    catching specified exceptions and providing feedback.

    :param func: Function to execute repeatedly.
    :param args: Positional arguments to pass to the function.
    :param exceptions: A single exception type or a tuple of exception types to catch and retry.
                       Defaults to catching all generic exceptions (Exception).
    :param on_exception: Action to take when a caught error occurs:
                     - str: The string message to print.
                     - Callable: A function (Callable[[BaseException], Any]) to call with the exception object.
                     - None (default): Prints the generic error message (f'Error: {e}').
    :param kwargs: Keyword arguments to pass to the function.
    :raises KeyboardInterrupt: Handled internally; terminates the program gracefully (exit 0).
    :return: The result returned by the successfully executed function.
    """
    while True:
        try:
            return func(*args, **kwargs)
        except KeyboardInterrupt:
            exit(0)
        except exceptions as e:
            if isinstance(on_exception, str):
                print(on_exception)
            elif callable(on_exception):
                on_exception(e)
            else:
                print(f'Error: {e}')
