import base64
import logging
import os
import traceback
from typing import Any, Callable, Dict, Optional, Union

import ray
import ray._private.ray_constants as ray_constants
import ray.cloudpickle as pickle
from ray._common.utils import load_class
from ray._private.function_manager import build_setup_hook_export_entry
from ray.runtime_env import RuntimeEnv

logger = logging.getLogger(__name__)

RUNTIME_ENV_FUNC_IDENTIFIER = "ray_runtime_env_func::"


def get_import_export_timeout():
    return int(
        os.environ.get(
            ray_constants.RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR, "60"
        )
    )


def decode_function_key(key: bytes) -> str:
    # b64encode only includes A-Z, a-z, 0-9, + and / characters
    return RUNTIME_ENV_FUNC_IDENTIFIER + base64.b64encode(key).decode()


def _encode_function_key(key: str) -> bytes:
    assert key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER)
    return base64.b64decode(key[len(RUNTIME_ENV_FUNC_IDENTIFIER) :])


def _raise_setup_hook_conflict(existing_hook_value: str, setup_hook_desc: str) -> None:
    raise RuntimeError(
        "Conflicting worker_process_setup_hook: the setup hook env "
        f"var is already set to '{existing_hook_value}', but "
        f"runtime_env specifies {setup_hook_desc}."
    )


def export_setup_func_callable(
    runtime_env: Union[Dict[str, Any], RuntimeEnv],
    setup_func: Callable,
    worker: "ray.Worker",
) -> Union[Dict[str, Any], RuntimeEnv]:
    assert isinstance(setup_func, Callable)
    try:
        key = worker.function_actor_manager.export_setup_func(
            setup_func, timeout=get_import_export_timeout()
        )
    except Exception as e:
        raise ray.exceptions.RuntimeEnvSetupError(
            "Failed to export the setup function."
        ) from e
    env_vars = runtime_env.get("env_vars", {})
    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
        "is not permitted because it is reserved for the internal use."
    )
    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = decode_function_key(key)
    runtime_env["env_vars"] = env_vars
    # Note: This field is no-op. We don't have a plugin for the setup hook
    # because we can implement it simply using an env var.
    # This field is just for the observability purpose, so we store
    # the name of the method.
    runtime_env["worker_process_setup_hook"] = setup_func.__name__
    return runtime_env


def export_setup_func_module(
    runtime_env: Union[Dict[str, Any], RuntimeEnv],
    setup_func_module: str,
) -> Union[Dict[str, Any], RuntimeEnv]:
    assert isinstance(setup_func_module, str)
    env_vars = runtime_env.get("env_vars", {})
    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
        "is not permitted because it is reserved for the internal use."
    )
    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = setup_func_module
    runtime_env["env_vars"] = env_vars
    return runtime_env


def _check_setup_hook_consistency(
    existing_hook_value: str,
    setup_func: Union[Callable, str],
    worker: "ray.Worker",
) -> None:
    """Validate that an already-set hook env var is consistent with setup_func.

    When the env var is already populated (e.g. inherited from a job supervisor),
    we compare it against the `worker_process_setup_hook` field in the runtime_env
    to detect silent mismatches.

    Args:
        existing_hook_value: The value of the existing hook env var.
        setup_func: The setup function or module path.
        worker: The worker instance.

    Raises:
        RuntimeError: If a conflict between the existing env var and setup_func is detected.
    """
    if isinstance(setup_func, Callable):
        try:
            _encode_function_key(existing_hook_value)
        except Exception:
            _raise_setup_hook_conflict(
                existing_hook_value, f"callable '{setup_func.__name__}'"
            )
        _check_callable_hooks_match(existing_hook_value, setup_func, worker)
    elif isinstance(setup_func, str):
        try:
            _encode_function_key(existing_hook_value)
            existing_is_callable_ref = True
        except Exception:
            existing_is_callable_ref = False
        if existing_is_callable_ref or existing_hook_value != setup_func:
            _raise_setup_hook_conflict(existing_hook_value, f"'{setup_func}'")


def _check_callable_hooks_match(
    existing_hook_value: str,
    setup_func: Callable,
    worker: "ray.Worker",
) -> None:
    """Verify a callable produces the same GCS key as the existing env var."""
    _, _, expected_key = build_setup_hook_export_entry(
        setup_func, worker.current_job_id.binary()
    )
    expected_env_value = decode_function_key(expected_key)

    if existing_hook_value != expected_env_value:
        _raise_setup_hook_conflict(
            existing_hook_value, f"callable '{setup_func.__name__}'"
        )


def upload_worker_process_setup_hook_if_needed(
    runtime_env: Union[Dict[str, Any], RuntimeEnv],
    worker: "ray.Worker",
) -> Union[Dict[str, Any], RuntimeEnv]:
    """Uploads the worker_process_setup_hook to GCS with a key.

    runtime_env["worker_process_setup_hook"] is converted to a decoded key
    that can load the worker setup hook function from GCS.
    i.e., you can use internalKV.Get(runtime_env["worker_process_setup_hook])
    to access the worker setup hook from GCS.

    Args:
        runtime_env: The runtime_env. The value will be modified
            when returned.
        worker: ray.worker instance.

    Returns:
        The modified runtime_env with the setup hook processed into an env var.
    """
    setup_func = runtime_env.get("worker_process_setup_hook")

    if setup_func is None:
        return runtime_env

    env_vars = runtime_env.get("env_vars", {})
    existing_hook = env_vars.get(ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR)

    if existing_hook is not None:
        # A setup hook is already populated (e.g. inherited from job supervisor).
        # Validate that it is consistent with the current worker_process_setup_hook.
        _check_setup_hook_consistency(existing_hook, setup_func, worker)
        return runtime_env

    if isinstance(setup_func, Callable):
        return export_setup_func_callable(runtime_env, setup_func, worker)
    elif isinstance(setup_func, str):
        return export_setup_func_module(runtime_env, setup_func)
    else:
        raise TypeError(
            "worker_process_setup_hook must be a function, " f"got {type(setup_func)}."
        )


def load_and_execute_setup_hook(
    worker_process_setup_hook_key: str,
) -> Optional[str]:
    """Load the setup hook from a given key and execute.

    Args:
        worker_process_setup_hook_key: The key to import the setup hook
            from GCS.
    Returns:
        An error message if it fails. None if it succeeds.
    """
    assert worker_process_setup_hook_key is not None
    if not worker_process_setup_hook_key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER):
        return load_and_execute_setup_hook_module(worker_process_setup_hook_key)
    else:
        return load_and_execute_setup_hook_func(worker_process_setup_hook_key)


def load_and_execute_setup_hook_module(
    worker_process_setup_hook_key: str,
) -> Optional[str]:
    try:
        setup_func = load_class(worker_process_setup_hook_key)
        setup_func()
        return None
    except Exception:
        error_message = (
            "Failed to execute the setup hook method, "
            f"{worker_process_setup_hook_key} "
            "from ``ray.init(runtime_env="
            f"{{'worker_process_setup_hook': {worker_process_setup_hook_key}}})``. "
            "Please make sure the given module exists and is available "
            "from ray workers. For more details, see the error trace below.\n"
            f"{traceback.format_exc()}"
        )
        return error_message


def load_and_execute_setup_hook_func(
    worker_process_setup_hook_key: str,
) -> Optional[str]:
    worker = ray._private.worker.global_worker
    assert worker.connected
    func_manager = worker.function_actor_manager
    try:
        worker_setup_func_info = func_manager.fetch_registered_method(
            _encode_function_key(worker_process_setup_hook_key),
            timeout=get_import_export_timeout(),
        )
    except Exception:
        error_message = (
            "Failed to import setup hook within "
            f"{get_import_export_timeout()} seconds.\n"
            f"{traceback.format_exc()}"
        )
        return error_message

    try:
        setup_func = pickle.loads(worker_setup_func_info.function)
    except Exception:
        error_message = (
            "Failed to deserialize the setup hook method.\n" f"{traceback.format_exc()}"
        )
        return error_message

    try:
        setup_func()
    except Exception:
        error_message = (
            f"Failed to execute the setup hook method. Function name:"
            f"{worker_setup_func_info.function_name}\n"
            f"{traceback.format_exc()}"
        )
        return error_message

    return None