import asyncio
import hashlib
import json
import logging
import os
import shutil
import sys
import tempfile
from typing import Dict, List, Optional, Tuple
from contextlib import asynccontextmanager
from asyncio import create_task, get_running_loop

from ray._private.runtime_env.context import RuntimeEnvContext
from ray._private.runtime_env.packaging import Protocol, parse_uri
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
from ray._private.runtime_env.utils import check_output_cmd
from ray._private.utils import get_directory_size_bytes, try_to_create_directory

default_logger = logging.getLogger(__name__)

_WIN32 = os.name == "nt"

INTERNAL_PIP_FILENAME = "ray_runtime_env_internal_pip_requirements.txt"
MAX_INTERNAL_PIP_FILENAME_TRIES = 100


def _get_pip_hash(pip_dict: Dict) -> str:
    serialized_pip_spec = json.dumps(pip_dict, sort_keys=True)
    hash = hashlib.sha1(serialized_pip_spec.encode("utf-8")).hexdigest()
    return hash


def get_uri(runtime_env: Dict) -> Optional[str]:
    """Return `"pip://<hashed_dependencies>"`, or None if no GC required."""
    pip = runtime_env.get("pip")
    if pip is not None:
        if isinstance(pip, dict):
            uri = "pip://" + _get_pip_hash(pip_dict=pip)
        elif isinstance(pip, list):
            uri = "pip://" + _get_pip_hash(pip_dict=dict(packages=pip))
        else:
            raise TypeError(
                "pip field received by RuntimeEnvAgent must be "
                f"list or dict, not {type(pip).__name__}."
            )
    else:
        uri = None
    return uri


class _PathHelper:
    @staticmethod
    def get_virtualenv_path(target_dir: str) -> str:
        return os.path.join(target_dir, "virtualenv")

    @classmethod
    def get_virtualenv_python(cls, target_dir: str) -> str:
        virtualenv_path = cls.get_virtualenv_path(target_dir)
        if _WIN32:
            return os.path.join(virtualenv_path, "Scripts", "python.exe")
        else:
            return os.path.join(virtualenv_path, "bin", "python")

    @classmethod
    def get_virtualenv_activate_command(cls, target_dir: str) -> List[str]:
        virtualenv_path = cls.get_virtualenv_path(target_dir)
        if _WIN32:
            cmd = [os.path.join(virtualenv_path, "Scripts", "activate.bat")]

        else:
            cmd = ["source", os.path.join(virtualenv_path, "bin/activate")]
        return cmd + ["1>&2", "&&"]

    @staticmethod
    def get_requirements_file(target_dir: str, pip_list: Optional[List[str]]) -> str:
        """Returns the path to the requirements file to use for this runtime env.

        If pip_list is not None, we will check if the internal pip filename is in any of
        the entries of pip_list. If so, we will append numbers to the end of the
        filename until we find one that doesn't conflict. This prevents infinite
        recursion if the user specifies the internal pip filename in their pip list.

        Args:
            target_dir: The directory to store the requirements file in.
            pip_list: A list of pip requirements specified by the user.

        Returns:
            The path to the requirements file to use for this runtime env.
        """

        def filename_in_pip_list(filename: str) -> bool:
            for pip_entry in pip_list:
                if filename in pip_entry:
                    return True
            return False

        filename = INTERNAL_PIP_FILENAME
        if pip_list is not None:
            i = 1
            while (
                filename_in_pip_list(filename) and i < MAX_INTERNAL_PIP_FILENAME_TRIES
            ):
                filename = f"{INTERNAL_PIP_FILENAME}.{i}"
                i += 1
            if i == MAX_INTERNAL_PIP_FILENAME_TRIES:
                raise RuntimeError(
                    "Could not find a valid filename for the internal "
                    "pip requirements file. Please specify a different "
                    "pip list in your runtime env."
                )
        return os.path.join(target_dir, filename)


class PipProcessor:
    def __init__(
        self,
        target_dir: str,
        runtime_env: "RuntimeEnv",  # noqa: F821
        logger: Optional[logging.Logger] = default_logger,
    ):
        try:
            import virtualenv  # noqa: F401 ensure virtualenv exists.
        except ImportError:
            raise RuntimeError(
                f"Please install virtualenv "
                f"`{sys.executable} -m pip install virtualenv`"
                f"to enable pip runtime env."
            )
        logger.debug("Setting up pip for runtime_env: %s", runtime_env)
        self._target_dir = target_dir
        self._runtime_env = runtime_env
        self._logger = logger

        self._pip_config = self._runtime_env.pip_config()
        self._pip_env = os.environ.copy()
        self._pip_env.update(self._runtime_env.env_vars())

    @staticmethod
    def _is_in_virtualenv() -> bool:
        # virtualenv <= 16.7.9 sets the real_prefix,
        # virtualenv > 16.7.9 & venv set the base_prefix.
        # So, we check both of them here.
        # https://github.com/pypa/virtualenv/issues/1622#issuecomment-586186094
        return hasattr(sys, "real_prefix") or (
            hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
        )

    @classmethod
    async def _ensure_pip_version(
        cls,
        path: str,
        pip_version: Optional[str],
        cwd: str,
        pip_env: Dict,
        logger: logging.Logger,
    ):
        """Run the pip command to reinstall pip to the specified version."""
        if not pip_version:
            return

        python = _PathHelper.get_virtualenv_python(path)
        # Ensure pip version.
        pip_reinstall_cmd = [
            python,
            "-m",
            "pip",
            "install",
            "--disable-pip-version-check",
            f"pip{pip_version}",
        ]
        logger.info("Installing pip with version %s", pip_version)

        await check_output_cmd(pip_reinstall_cmd, logger=logger, cwd=cwd, env=pip_env)

    async def _pip_check(
        self,
        path: str,
        pip_check: bool,
        cwd: str,
        pip_env: Dict,
        logger: logging.Logger,
    ):
        """Run the pip check command to check python dependency conflicts.
        If exists conflicts, the exit code of pip check command will be non-zero.
        """
        if not pip_check:
            logger.info("Skip pip check.")
            return
        python = _PathHelper.get_virtualenv_python(path)

        await check_output_cmd(
            [python, "-m", "pip", "check", "--disable-pip-version-check"],
            logger=logger,
            cwd=cwd,
            env=pip_env,
        )

        logger.info("Pip check on %s successfully.", path)

    @staticmethod
    @asynccontextmanager
    async def _check_ray(python: str, cwd: str, logger: logging.Logger):
        """A context manager to check ray is not overwritten.

        Currently, we only check ray version and path. It works for virtualenv,
          - ray is in Python's site-packages.
          - ray is overwritten during yield.
          - ray is in virtualenv's site-packages.
        """

        async def _get_ray_version_and_path() -> Tuple[str, str]:
            with tempfile.TemporaryDirectory(
                prefix="check_ray_version_tempfile"
            ) as tmp_dir:
                ray_version_path = os.path.join(tmp_dir, "ray_version.txt")
                check_ray_cmd = [
                    python,
                    "-c",
                    """
import ray
with open(r"{ray_version_path}", "wt") as f:
    f.write(ray.__version__)
    f.write(" ")
    f.write(ray.__path__[0])
                    """.format(
                        ray_version_path=ray_version_path
                    ),
                ]
                if _WIN32:
                    env = os.environ.copy()
                else:
                    env = {}
                output = await check_output_cmd(
                    check_ray_cmd, logger=logger, cwd=cwd, env=env
                )
                logger.info(
                    f"try to write ray version information in: {ray_version_path}"
                )
                with open(ray_version_path, "rt") as f:
                    output = f.read()
                # print after import ray may have [0m endings, so we strip them by *_
                ray_version, ray_path, *_ = [s.strip() for s in output.split()]
            return ray_version, ray_path

        version, path = await _get_ray_version_and_path()
        yield
        actual_version, actual_path = await _get_ray_version_and_path()
        if actual_version != version or actual_path != path:
            raise RuntimeError(
                "Changing the ray version is not allowed: \n"
                f"  current version: {actual_version}, "
                f"current path: {actual_path}\n"
                f"  expect version: {version}, "
                f"expect path: {path}\n"
                "Please ensure the dependencies in the runtime_env pip field "
                "do not install a different version of Ray."
            )

    @classmethod
    async def _create_or_get_virtualenv(
        cls, path: str, cwd: str, logger: logging.Logger
    ):
        """Create or get a virtualenv from path."""
        python = sys.executable
        virtualenv_path = os.path.join(path, "virtualenv")
        virtualenv_app_data_path = os.path.join(path, "virtualenv_app_data")

        if _WIN32:
            current_python_dir = sys.prefix
            env = os.environ.copy()
        else:
            current_python_dir = os.path.abspath(
                os.path.join(os.path.dirname(python), "..")
            )
            env = {}

        if cls._is_in_virtualenv():
            # virtualenv-clone homepage:
            # https://github.com/edwardgeorge/virtualenv-clone
            # virtualenv-clone Usage:
            # virtualenv-clone /path/to/existing/venv /path/to/cloned/ven
            # or
            # python -m clonevirtualenv /path/to/existing/venv /path/to/cloned/ven
            clonevirtualenv = os.path.join(
                os.path.dirname(__file__), "_clonevirtualenv.py"
            )
            create_venv_cmd = [
                python,
                clonevirtualenv,
                current_python_dir,
                virtualenv_path,
            ]
            logger.info(
                "Cloning virtualenv %s to %s", current_python_dir, virtualenv_path
            )
        else:
            # virtualenv options:
            # https://virtualenv.pypa.io/en/latest/cli_interface.html
            #
            # --app-data
            # --reset-app-data
            #   Set an empty seperated app data folder for current virtualenv.
            #
            # --no-periodic-update
            #   Disable the periodic (once every 14 days) update of the embedded
            #   wheels.
            #
            # --system-site-packages
            #   Inherit site packages.
            #
            # --no-download
            #   Never download the latest pip/setuptools/wheel from PyPI.
            create_venv_cmd = [
                python,
                "-m",
                "virtualenv",
                "--app-data",
                virtualenv_app_data_path,
                "--reset-app-data",
                "--no-periodic-update",
                "--system-site-packages",
                "--no-download",
                virtualenv_path,
            ]
            logger.info(
                "Creating virtualenv at %s, current python dir %s",
                virtualenv_path,
                virtualenv_path,
            )
        await check_output_cmd(create_venv_cmd, logger=logger, cwd=cwd, env=env)

    @classmethod
    async def _install_pip_packages(
        cls,
        path: str,
        pip_packages: List[str],
        cwd: str,
        pip_env: Dict,
        logger: logging.Logger,
    ):
        virtualenv_path = _PathHelper.get_virtualenv_path(path)
        python = _PathHelper.get_virtualenv_python(path)
        # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ...
        pip_requirements_file = _PathHelper.get_requirements_file(path, pip_packages)

        def _gen_requirements_txt():
            with open(pip_requirements_file, "w") as file:
                for line in pip_packages:
                    file.write(line + "\n")

        # Avoid blocking the event loop.
        loop = get_running_loop()
        await loop.run_in_executor(None, _gen_requirements_txt)

        # pip options
        #
        # --disable-pip-version-check
        #   Don't periodically check PyPI to determine whether a new version
        #   of pip is available for download.
        #
        # --no-cache-dir
        #   Disable the cache, the pip runtime env is a one-time installation,
        #   and we don't need to handle the pip cache broken.
        pip_install_cmd = [
            python,
            "-m",
            "pip",
            "install",
            "--disable-pip-version-check",
            "--no-cache-dir",
            "-r",
            pip_requirements_file,
        ]
        logger.info("Installing python requirements to %s", virtualenv_path)

        await check_output_cmd(pip_install_cmd, logger=logger, cwd=cwd, env=pip_env)

    async def _run(self):
        path = self._target_dir
        logger = self._logger
        pip_packages = self._pip_config["packages"]
        # We create an empty directory for exec cmd so that the cmd will
        # run more stable. e.g. if cwd has ray, then checking ray will
        # look up ray in cwd instead of site packages.
        exec_cwd = os.path.join(path, "exec_cwd")
        os.makedirs(exec_cwd, exist_ok=True)
        try:
            await self._create_or_get_virtualenv(path, exec_cwd, logger)
            python = _PathHelper.get_virtualenv_python(path)
            async with self._check_ray(python, exec_cwd, logger):
                # Ensure pip version.
                await self._ensure_pip_version(
                    path,
                    self._pip_config.get("pip_version", None),
                    exec_cwd,
                    self._pip_env,
                    logger,
                )
                # Install pip packages.
                await self._install_pip_packages(
                    path,
                    pip_packages,
                    exec_cwd,
                    self._pip_env,
                    logger,
                )
                # Check python environment for conflicts.
                await self._pip_check(
                    path,
                    self._pip_config.get("pip_check", False),
                    exec_cwd,
                    self._pip_env,
                    logger,
                )
        except Exception:
            logger.info("Delete incomplete virtualenv: %s", path)
            shutil.rmtree(path, ignore_errors=True)
            logger.exception("Failed to install pip packages.")
            raise

    def __await__(self):
        return self._run().__await__()


class PipPlugin(RuntimeEnvPlugin):
    name = "pip"

    def __init__(self, resources_dir: str):
        self._pip_resources_dir = os.path.join(resources_dir, "pip")
        self._creating_task = {}
        # Maps a URI to a lock that is used to prevent multiple concurrent
        # installs of the same virtualenv, see #24513
        self._create_locks: Dict[str, asyncio.Lock] = {}
        try_to_create_directory(self._pip_resources_dir)

    def _get_path_from_hash(self, hash: str) -> str:
        """Generate a path from the hash of a pip spec.

        Example output:
            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
                /pip/ray-9a7972c3a75f55e976e620484f58410c920db091
        """
        return os.path.join(self._pip_resources_dir, hash)

    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
        """Return the pip URI from the RuntimeEnv if it exists, else return []."""
        pip_uri = runtime_env.pip_uri()
        if pip_uri:
            return [pip_uri]
        return []

    def delete_uri(
        self, uri: str, logger: Optional[logging.Logger] = default_logger
    ) -> int:
        """Delete URI and return the number of bytes deleted."""
        logger.info("Got request to delete pip URI %s", uri)
        protocol, hash = parse_uri(uri)
        if protocol != Protocol.PIP:
            raise ValueError(
                "PipPlugin can only delete URIs with protocol "
                f"pip. Received protocol {protocol}, URI {uri}"
            )

        # Cancel running create task.
        task = self._creating_task.pop(hash, None)
        if task is not None:
            task.cancel()

        pip_env_path = self._get_path_from_hash(hash)
        local_dir_size = get_directory_size_bytes(pip_env_path)
        del self._create_locks[uri]
        try:
            shutil.rmtree(pip_env_path)
        except OSError as e:
            logger.warning(f"Error when deleting pip env {pip_env_path}: {str(e)}")
            return 0

        return local_dir_size

    async def create(
        self,
        uri: str,
        runtime_env: "RuntimeEnv",  # noqa: F821
        context: RuntimeEnvContext,
        logger: Optional[logging.Logger] = default_logger,
    ) -> int:
        if not runtime_env.has_pip():
            return 0

        protocol, hash = parse_uri(uri)
        target_dir = self._get_path_from_hash(hash)

        async def _create_for_hash():
            await PipProcessor(
                target_dir,
                runtime_env,
                logger,
            )

            loop = get_running_loop()
            return await loop.run_in_executor(
                None, get_directory_size_bytes, target_dir
            )

        if uri not in self._create_locks:
            # async lock to prevent the same virtualenv being concurrently installed
            self._create_locks[uri] = asyncio.Lock()

        async with self._create_locks[uri]:
            self._creating_task[hash] = task = create_task(_create_for_hash())
            task.add_done_callback(lambda _: self._creating_task.pop(hash, None))
            return await task

    def modify_context(
        self,
        uris: List[str],
        runtime_env: "RuntimeEnv",  # noqa: F821
        context: RuntimeEnvContext,
        logger: logging.Logger = default_logger,
    ):
        if not runtime_env.has_pip():
            return
        # PipPlugin only uses a single URI.
        uri = uris[0]
        # Update py_executable.
        protocol, hash = parse_uri(uri)
        target_dir = self._get_path_from_hash(hash)
        virtualenv_python = _PathHelper.get_virtualenv_python(target_dir)

        if not os.path.exists(virtualenv_python):
            raise ValueError(
                f"Local directory {target_dir} for URI {uri} does "
                "not exist on the cluster. Something may have gone wrong while "
                "installing the runtime_env `pip` packages."
            )
        context.py_executable = virtualenv_python
        context.command_prefix += _PathHelper.get_virtualenv_activate_command(
            target_dir
        )
