import asyncio
import contextlib
import functools
import logging
import threading
import time
import traceback
from datetime import datetime
from enum import Enum
from typing import (
    Any,
    Callable,
    ContextManager,
    Dict,
    Generator,
    Generic,
    Iterator,
    List,
    Optional,
    TypeVar,
    Union,
)

import ray
from ray.train._internal.utils import count_required_parameters
from ray.train.v2._internal.exceptions import UserExceptionWithTraceback

logger = logging.getLogger(__name__)


T = TypeVar("T")


def bundle_to_remote_args(bundle: dict) -> dict:
    """Convert a bundle of resources to Ray actor/task arguments.

    >>> bundle_to_remote_args({"GPU": 1, "memory": 1, "custom": 0.1})
    {'num_cpus': 0, 'num_gpus': 1, 'memory': 1, 'resources': {'custom': 0.1}}
    """
    bundle = bundle.copy()
    args = {
        "num_cpus": bundle.pop("CPU", 0),
        "num_gpus": bundle.pop("GPU", 0),
        "memory": bundle.pop("memory", 0),
    }
    if bundle:
        args["resources"] = bundle
    return args


def construct_train_func(
    train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]],
    config: Optional[Dict[str, Any]],
    train_func_context: ContextManager,
    fn_arg_name: Optional[str] = "train_loop_per_worker",
) -> Callable[[], T]:
    """Validates and constructs the training function to execute.

    Args:
        train_func: The training function to execute.
            This can either take in no arguments or a ``config`` dict.
        config: Configurations to pass into ``train_func``. If None then an
            empty Dict will be created.
        train_func_context: Context manager for user's `train_func`, which executes
            backend-specific logic before and after the training function.
        fn_arg_name: The name of training function to use for error messages.

    Returns:
        A valid training function.

    Raises:
        ValueError: if the input ``train_func`` is invalid.
    """
    num_required_params = count_required_parameters(train_func)

    if num_required_params > 1:
        err_msg = (
            f"{fn_arg_name} should take in 0 or 1 required arguments, but it accepts "
            f"{num_required_params} required arguments instead."
        )
        raise ValueError(err_msg)

    if num_required_params == 1:
        config = config or {}

        @functools.wraps(train_func)
        def train_fn():
            with train_func_context():
                return train_func(config)

    else:  # num_params == 0

        @functools.wraps(train_func)
        def train_fn():
            with train_func_context():
                return train_func()

    return train_fn


class TrainingFramework(Enum):
    TORCH = "torch"
    JAX = "jax"
    TENSORFLOW = "tensorflow"
    XGBOOST = "xgboost"
    LIGHTGBM = "lightgbm"

    def module_names(self) -> tuple[str, ...]:
        """Returns the relevant module names for the training framework.

        These module names are used by Train state version collection (see
        `_get_framework_version`) to gather versions of key framework-related packages.

        Note: If adding a new module, make sure to use the module name rather than
        the distribution name. (e.g. sklearn instead of scikit-learn)
        """
        if self is TrainingFramework.TORCH:
            return ("torch",)
        if self is TrainingFramework.JAX:
            return ("jax", "jaxlib")
        if self is TrainingFramework.TENSORFLOW:
            return ("tensorflow", "keras")
        if self is TrainingFramework.XGBOOST:
            return ("xgboost",)
        if self is TrainingFramework.LIGHTGBM:
            return ("lightgbm",)

        return (self.value,)


class ObjectRefWrapper(Generic[T]):
    """Thin wrapper around ray.put to manually control dereferencing."""

    def __init__(self, obj: T):
        self._ref = ray.put(obj)

    def get(self) -> T:
        return ray.get(self._ref)


def date_str(include_ms: bool = False):
    pattern = "%Y-%m-%d_%H-%M-%S"
    if include_ms:
        pattern += ".%f"
    return datetime.today().strftime(pattern)


def time_monotonic():
    return time.monotonic()


def _copy_doc(copy_func):
    def wrapped(func):
        func.__doc__ = copy_func.__doc__
        return func

    return wrapped


@contextlib.contextmanager
def invoke_context_managers(
    context_managers: List[ContextManager],
) -> Generator[None, None, None]:
    """
    Utility to invoke a list of context managers and yield sequentially.

    Args:
        context_managers: List of context managers to invoke.
    """
    with contextlib.ExitStack() as stack:
        for context_manager in context_managers:
            stack.enter_context(context_manager())
        yield


def get_module_name(obj: object) -> str:
    """Returns the full module name of the given object, including its qualified name.

    Args:
        obj: The object (class, function, etc.) whose module name is required.

    Returns:
        Full module and qualified name as a string.
    """
    return f"{obj.__module__}.{obj.__qualname__}"


def get_callable_name(fn: Callable) -> str:
    """Returns a readable name for any callable.

    Args:
        fn: The callable to extract a name from.

    Returns:
        A human-readable name for the callable.

    Examples:

        >>> get_callable_name(lambda x: x)
        '<lambda>'
        >>> def foo(a, b): pass
        >>> get_callable_name(foo)
        'foo'
        >>> from functools import partial
        >>> bar = partial(partial(foo, a=1), b=2)
        >>> get_callable_name(bar)
        'foo'
        >>> class Dummy:
        ...     def __call__(self, a, b): pass
        >>> get_callable_name(Dummy())
        'Dummy'
    """
    if isinstance(fn, functools.partial):
        return get_callable_name(fn.func)

    # Use __name__ for regular functions and lambdas
    if hasattr(fn, "__name__"):
        return fn.__name__

    # Fallback to the class name for objects that implement __call__
    return fn.__class__.__name__


def construct_user_exception_with_traceback(
    e: BaseException, exclude_frames: int = 0
) -> UserExceptionWithTraceback:
    """Construct a UserExceptionWithTraceback from a base exception.

    Args:
        e: The base exception to construct a UserExceptionWithTraceback from.
        exclude_frames: The number of frames to exclude from the beginnning of
            the traceback.

    Returns:
        A UserExceptionWithTraceback object.
    """
    # TODO(justinvyu): This is brittle and may break if the call stack
    # changes. Figure out a more robust way to exclude these frames.
    exc_traceback_str = traceback.format_exc(
        limit=-(len(traceback.extract_tb(e.__traceback__)) - exclude_frames)
    )
    logger.error(f"Error in training function:\n{exc_traceback_str}")
    return UserExceptionWithTraceback(e, traceback_str=exc_traceback_str)


def _in_ray_train_worker() -> bool:
    """Check if the current process is a Ray Train V2 worker."""
    from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils

    try:
        get_train_fn_utils()
        return True
    except RuntimeError:
        return False


def requires_train_worker(raise_in_tune_session: bool = False) -> Callable:
    """Check that the caller is a Ray Train worker spawned by Ray Train,
    with access to training function utilities.

    Args:
        raise_in_tune_session: Whether to raise a specific error message if the caller
            is in a Tune session. If True, will raise a DeprecationWarning.

    Returns:
        A decorator that performs this check, which raises an error if the caller
        is not a Ray Train worker.
    """

    def _wrap(fn: Callable) -> Callable:
        @functools.wraps(fn)
        def _wrapped_fn(*args, **kwargs):
            from ray.tune.trainable.trainable_fn_utils import _in_tune_session

            if raise_in_tune_session and _in_tune_session():
                raise DeprecationWarning(
                    f"`ray.train.{fn.__name__}` is deprecated when running in a function "
                    "passed to Ray Tune. Please use the equivalent `ray.tune` API instead. "
                    "See this issue for more context: "
                    "https://github.com/ray-project/ray/issues/49454"
                )

            if not _in_ray_train_worker():
                raise RuntimeError(
                    f"`{fn.__name__}` cannot be used outside of a Ray Train training function. "
                    "You are calling this API from the driver or another non-training process. "
                    "These utilities are only available within a function launched by `trainer.fit()`."
                )
            return fn(*args, **kwargs)

        return _wrapped_fn

    return _wrap


async def wait_with_logging(
    condition: asyncio.Condition,
    predicate: Optional[Callable[[], bool]] = None,
    generate_warning_message: Optional[Callable[[], str]] = None,
    warn_interval_s: float = 60,
    timeout_s: Optional[float] = None,
):
    """Waits for condition to be notified, logging warnings and eventually timing out.

    You must acquire the condition before calling this function.

    Args:
        condition: The condition to wait for.
        predicate: Wait until this predicate is True. If None, wait until the condition
            is notified.
        generate_warning_message: A function that generates the warning message to log.
            If None, no warning is logged.
        warn_interval_s: The interval in seconds to log a warning.
        timeout_s: The timeout in seconds. Defaults to``None`` to not time out.
    """

    async def _wait_loop():
        while True:
            try:
                await asyncio.wait_for(
                    condition.wait()
                    if predicate is None
                    else condition.wait_for(predicate),
                    timeout=warn_interval_s,
                )
                return
            # asyncio.wait_for() raises `asyncio.TimeoutError` for asyncio<=3.10
            # and raises `TimeoutError` for asyncio>=3.11
            # https://docs.python.org/3/library/asyncio-task.html#asyncio.wait_for
            except (asyncio.TimeoutError, TimeoutError):
                if generate_warning_message is not None:
                    warning_message = generate_warning_message()
                    logger.warning(warning_message)

    await asyncio.wait_for(
        _wait_loop(),
        timeout=timeout_s,
    )


@contextlib.contextmanager
def context_watchdog(fn: Callable, *args: Any) -> Iterator[None]:
    """Run a function in a background thread for the duration of the context.

    The function is started in a daemon thread on entry. On exit, a
    threading.Event is set to signal the thread to stop. The function is
    responsible for checking the event and returning promptly once it is set.

    Args:
        fn: A function whose first argument is a threading.Event stop signal.
            The function should return when stop_event.is_set() or
            stop_event.wait(...) returns True.
        *args: Additional arguments forwarded to fn after the stop event.

    Yields:
        None: Control is yielded to the caller while the watchdog thread runs.
    """
    stop_event = threading.Event()
    thread = threading.Thread(
        target=fn,
        args=(stop_event, *args),
        daemon=True,  # thread will end even if the finally is bypassed by an abnormal exit
    )
    thread.start()
    try:
        yield
    finally:
        stop_event.set()
        thread.join()