import os
import threading
from contextlib import contextmanager
from functools import wraps
from typing import Any, Callable, Optional, Tuple, TypeVar, cast, overload

from ray._private.auto_init_hook import auto_init_ray

F = TypeVar("F", bound=Callable[..., Any])

# Attr set on func defs to mark they have been converted to client mode.
RAY_CLIENT_MODE_ATTR = "__ray_client_mode_key__"

# Global setting of whether client mode is enabled. This default to OFF,
# but is enabled upon ray.client(...).connect() or in tests.
is_client_mode_enabled = os.environ.get("RAY_CLIENT_MODE", "0") == "1"

# When RAY_CLIENT_MODE == 1, we treat it as default enabled client mode
# This is useful for testing
is_client_mode_enabled_by_default = is_client_mode_enabled
os.environ.update({"RAY_CLIENT_MODE": "0"})

is_init_called = False

# Local setting of whether to ignore client hook conversion. This defaults
# to TRUE and is disabled when the underlying 'real' Ray function is needed.
_client_hook_status_on_thread = threading.local()
_client_hook_status_on_thread.status = True


def _get_client_hook_status_on_thread():
    """Get's the value of `_client_hook_status_on_thread`.
    Since `_client_hook_status_on_thread` is a thread-local variable, we may
    need to add and set the 'status' attribute.
    """
    global _client_hook_status_on_thread
    if not hasattr(_client_hook_status_on_thread, "status"):
        _client_hook_status_on_thread.status = True
    return _client_hook_status_on_thread.status


def _set_client_hook_status(val: bool):
    global _client_hook_status_on_thread
    _client_hook_status_on_thread.status = val


def _disable_client_hook():
    global _client_hook_status_on_thread
    out = _get_client_hook_status_on_thread()
    _client_hook_status_on_thread.status = False
    return out


def _explicitly_enable_client_mode():
    """Force client mode to be enabled.
    NOTE: This should not be used in tests, use `enable_client_mode`.
    """
    global is_client_mode_enabled
    is_client_mode_enabled = True


def _explicitly_disable_client_mode():
    global is_client_mode_enabled
    is_client_mode_enabled = False


@contextmanager
def disable_client_hook():
    val = _disable_client_hook()
    try:
        yield None
    finally:
        _set_client_hook_status(val)


@contextmanager
def enable_client_mode():
    _explicitly_enable_client_mode()
    try:
        yield None
    finally:
        _explicitly_disable_client_mode()


@overload
def client_mode_hook(func: F) -> F:
    ...


@overload
def client_mode_hook(*, local_only_kwargs: Tuple[str, ...] = ()) -> Callable[[F], F]:
    ...


def client_mode_hook(
    func: Optional[F] = None, *, local_only_kwargs: Tuple[str, ...] = ()
):
    """Decorator for whether to use the 'regular' ray version of a function,
    or the Ray Client version of that function.

    Args:
        func: This function. This is set when this function is used
            as a bare decorator.
        local_only_kwargs: Names of keyword arguments that apply only to the
            local (non-client) implementation. They are stripped before the
            call is redirected to the Ray Client, so the client API does not
            need to accept them. On the local path they pass through unchanged.
            as a decorator.

    Returns:
        The wrapped function that dispatches to the regular or client version.
    """

    def decorator(func: F) -> F:
        from ray.util.client import ray

        @wraps(func)
        def wrapper(*args, **kwargs):
            # NOTE(hchen): DO NOT use "import" inside this function.
            # Because when it's called within a `__del__` method, this error
            # will be raised (see #35114):
            # ImportError: sys.meta_path is None, Python is likely shutting down.
            if client_mode_should_convert():
                # Legacy code
                # we only convert init function if RAY_CLIENT_MODE=1
                if func.__name__ != "init" or is_client_mode_enabled_by_default:
                    if local_only_kwargs:
                        kwargs = {
                            k: v
                            for k, v in kwargs.items()
                            if k not in local_only_kwargs
                        }
                    return getattr(ray, func.__name__)(*args, **kwargs)
            return func(*args, **kwargs)

        return cast(F, wrapper)

    # Support both `@client_mode_hook` and
    # `@client_mode_hook(local_only_kwargs=...)` usage.
    if func is not None:
        return decorator(func)
    return decorator


def client_mode_should_convert():
    """Determines if functions should be converted to client mode."""

    # `is_client_mode_enabled_by_default` is used for testing with
    # `RAY_CLIENT_MODE=1`. This flag means all tests run with client mode.
    return (
        is_client_mode_enabled or is_client_mode_enabled_by_default
    ) and _get_client_hook_status_on_thread()


def client_mode_wrap(func):
    """Wraps a function called during client mode for execution as a remote
    task.

    Can be used to implement public features of ray client which do not
    belong in the main ray API (`ray.*`), yet require server-side execution.
    An example is the creation of placement groups:
    `ray.util.placement_group.placement_group()`. When called on the client
    side, this function is wrapped in a task to facilitate interaction with
    the GCS.
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        from ray.util.client import ray

        auto_init_ray()
        # Directly pass this through since `client_mode_wrap` is for
        # Placement Group APIs
        if client_mode_should_convert():
            f = ray.remote(num_cpus=0)(func)
            ref = f.remote(*args, **kwargs)
            return ray.get(ref)
        return func(*args, **kwargs)

    return wrapper


def client_mode_convert_function(func_cls, in_args, in_kwargs, **kwargs):
    """Runs a preregistered ray RemoteFunction through the ray client.

    The common case for this is to transparently convert that RemoteFunction
    to a ClientRemoteFunction. This happens in circumstances where the
    RemoteFunction is declared early, in a library and only then is Ray used in
    client mode -- necessitating a conversion.
    """
    from ray.util.client import ray

    key = getattr(func_cls, RAY_CLIENT_MODE_ATTR, None)

    # Second part of "or" is needed in case func_cls is reused between Ray
    # client sessions in one Python interpreter session.
    if (key is None) or (not ray._converted_key_exists(key)):
        key = ray._convert_function(func_cls)
        setattr(func_cls, RAY_CLIENT_MODE_ATTR, key)
    client_func = ray._get_converted(key)
    return client_func._remote(in_args, in_kwargs, **kwargs)


def client_mode_convert_actor(actor_cls, in_args, in_kwargs, **kwargs):
    """Runs a preregistered actor class on the ray client

    The common case for this decorator is for instantiating an ActorClass
    transparently as a ClientActorClass. This happens in circumstances where
    the ActorClass is declared early, in a library and only then is Ray used in
    client mode -- necessitating a conversion.
    """
    from ray.util.client import ray

    key = getattr(actor_cls, RAY_CLIENT_MODE_ATTR, None)
    # Second part of "or" is needed in case actor_cls is reused between Ray
    # client sessions in one Python interpreter session.
    if (key is None) or (not ray._converted_key_exists(key)):
        key = ray._convert_actor(actor_cls)
        setattr(actor_cls, RAY_CLIENT_MODE_ATTR, key)
    client_actor = ray._get_converted(key)
    return client_actor._remote(in_args, in_kwargs, **kwargs)