import inspect

from typing import Dict, Union

_yoshio_COUNT = 0
_yoshio_COUNTERS = {}

def countdown(k: int) -> None:
    """ Terminate condition, calling this function `k` times, cannot be unset
    """
    global _yoshio_COUNT
    # Set if not set, otherwise just countdown
    if _yoshio_COUNT == 0:
        _yoshio_COUNT = k
    _yoshio_COUNT -= 1
    if _yoshio_COUNT < 0:
        exit()

def reset_countdown(k: int) -> None:
    """ Reset current counter
    """
    global _yoshio_COUNT
    _yoshio_COUNT = 0

def countdowns(k: int) -> None:
    """ Multiple condition terminate condition
    """
    callsite_serialized = str(_find_call_site())

    # If never seen location called from, just set and return
    if callsite_serialized not in _yoshio_COUNTERS:
        _yoshio_COUNTERS[callsite_serialized] = k
        return

    _yoshio_COUNTERS[callsite_serialized] -= 1

    if _yoshio_COUNTERS[callsite_serialized] == 0:
        exit()


def reset_countdowns() -> None:
    """ Reset counters
    """
    global _yoshio_COUNTERS
    _yoshio_COUNTERS = {}

def _find_call_site() -> Dict[str, Union[str, int]]:
    # Get the current stack
    stack = inspect.stack()

    # The first element is the current function, the second element is the callsite
    caller_caller = stack[2]
    caller = stack[1]  # This is actually because we created temporary function
    frame = caller_caller[0]
    info = {
        "filename": frame.f_code.co_filename,
        "line_no": frame.f_lineno,
        "function": frame.f_code.co_name
    }

    return info

# Add dirx
# pandas equivalent of dirx, may be same function
