import hashlib
import json
import logging
import os
import shutil
import subprocess
from typing import List, Optional, Tuple, Union

"""Utilities for conda.  Adapted from https://github.com/mlflow/mlflow."""

# Name of environment variable indicating a path to a conda installation. Ray
# will default to running "conda" if unset.
RAY_CONDA_HOME = "RAY_CONDA_HOME"

_WIN32 = os.name == "nt"


def get_conda_activate_commands(conda_env_name: str) -> List[str]:
    """
    Get a list of commands to run to silently activate the given conda env.
    """
    #  Checking for newer conda versions
    if not _WIN32 and ("CONDA_EXE" in os.environ or RAY_CONDA_HOME in os.environ):
        conda_path = get_conda_bin_executable("conda")
        activate_conda_env = [
            ".",
            f"{os.path.dirname(conda_path)}/../etc/profile.d/conda.sh",
            "&&",
        ]
        activate_conda_env += ["conda", "activate", conda_env_name]

    else:
        activate_path = get_conda_bin_executable("activate")
        if not _WIN32:
            # Use bash command syntax
            activate_conda_env = ["source", activate_path, conda_env_name]
        else:
            conda_path = get_conda_bin_executable("conda")
            activate_conda_env = [conda_path, "activate", conda_env_name]
    return activate_conda_env + ["1>&2", "&&"]


def get_conda_bin_executable(executable_name: str) -> str:
    """
    Return path to the specified executable, assumed to be discoverable within
    a conda installation.

    The conda home directory (expected to contain a 'bin' subdirectory on
    linux) is configurable via the ``RAY_CONDA_HOME`` environment variable. If
    ``RAY_CONDA_HOME`` is unspecified, try the ``CONDA_EXE`` environment
    variable set by activating conda. If neither is specified, this method
    returns `executable_name`.
    """
    conda_home = os.environ.get(RAY_CONDA_HOME)
    if conda_home:
        if _WIN32:
            candidate = os.path.join(conda_home, "%s.exe" % executable_name)
            if os.path.exists(candidate):
                return candidate
            candidate = os.path.join(conda_home, "%s.bat" % executable_name)
            if os.path.exists(candidate):
                return candidate
        else:
            return os.path.join(conda_home, "bin/%s" % executable_name)
    else:
        conda_home = "."
    # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126
    if "CONDA_EXE" in os.environ:
        conda_bin_dir = os.path.dirname(os.environ["CONDA_EXE"])
        if _WIN32:
            candidate = os.path.join(conda_home, "%s.exe" % executable_name)
            if os.path.exists(candidate):
                return candidate
            candidate = os.path.join(conda_home, "%s.bat" % executable_name)
            if os.path.exists(candidate):
                return candidate
        else:
            return os.path.join(conda_bin_dir, executable_name)
    if _WIN32:
        return executable_name + ".bat"
    return executable_name


def _get_conda_env_name(conda_env_path: str) -> str:
    conda_env_contents = open(conda_env_path).read()
    return "ray-%s" % hashlib.sha1(conda_env_contents.encode("utf-8")).hexdigest()


def create_conda_env_if_needed(
    conda_yaml_file: str, prefix: str, logger: Optional[logging.Logger] = None
) -> None:
    """
    Given a conda YAML, creates a conda environment containing the required
    dependencies if such a conda environment doesn't already exist.
    Args:
        conda_yaml_file: The path to a conda `environment.yml` file.
        prefix: Directory to install the environment into via
            the `--prefix` option to conda create.  This also becomes the name
            of the conda env; i.e. it can be passed into `conda activate` and
            `conda remove`
        logger: Logger used to surface progress and errors; defaults to the
            module logger when not provided.
    """
    if logger is None:
        logger = logging.getLogger(__name__)

    conda_path = get_conda_bin_executable("conda")
    try:
        exec_cmd([conda_path, "--help"], throw_on_error=False)
    except (EnvironmentError, FileNotFoundError):
        raise ValueError(
            f"Could not find Conda executable at '{conda_path}'. "
            "Ensure Conda is installed as per the instructions at "
            "https://conda.io/projects/conda/en/latest/"
            "user-guide/install/index.html. "
            "You can also configure Ray to look for a specific "
            f"Conda executable by setting the {RAY_CONDA_HOME} "
            "environment variable to the path of the Conda executable."
        )

    _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"])
    envs = json.loads(stdout[stdout.index("{") :])["envs"]

    if prefix in envs:
        logger.info(f"Conda environment {prefix} already exists.")
        return

    create_cmd = [
        conda_path,
        "env",
        "create",
        "--file",
        conda_yaml_file,
        "--prefix",
        prefix,
    ]

    logger.info(f"Creating conda environment {prefix}")
    exit_code, output = exec_cmd_stream_to_logger(create_cmd, logger)
    if exit_code != 0:
        if os.path.exists(prefix):
            shutil.rmtree(prefix)
        raise RuntimeError(
            f"Failed to install conda environment {prefix}:\nOutput:\n{output}"
        )


def delete_conda_env(prefix: str, logger: Optional[logging.Logger] = None) -> bool:
    if logger is None:
        logger = logging.getLogger(__name__)

    logger.info(f"Deleting conda environment {prefix}")

    conda_path = get_conda_bin_executable("conda")
    delete_cmd = [conda_path, "remove", "-p", prefix, "--all", "-y"]
    exit_code, output = exec_cmd_stream_to_logger(delete_cmd, logger)

    if exit_code != 0:
        logger.debug(f"Failed to delete conda environment {prefix}:\n{output}")
        return False

    return True


def get_conda_env_list() -> list:
    """
    Get conda env list in full paths.
    """
    conda_path = get_conda_bin_executable("conda")
    try:
        exec_cmd([conda_path, "--help"], throw_on_error=False)
    except EnvironmentError:
        raise ValueError(f"Could not find Conda executable at {conda_path}.")
    _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"])
    envs = json.loads(stdout)["envs"]
    return envs


def get_conda_info_json() -> dict:
    """
    Get `conda info --json` output.

    Returns dict of conda info. See [1] for more details. We mostly care about these
    keys:

    - `conda_prefix`: str The path to the conda installation.
    - `envs`: List[str] absolute paths to conda environments.

    [1] https://github.com/conda/conda/blob/main/conda/cli/main_info.py
    """
    conda_path = get_conda_bin_executable("conda")
    try:
        exec_cmd([conda_path, "--help"], throw_on_error=False)
    except EnvironmentError:
        raise ValueError(f"Could not find Conda executable at {conda_path}.")
    _, stdout, _ = exec_cmd([conda_path, "info", "--json"])
    return json.loads(stdout)


def get_conda_envs(conda_info: dict) -> List[Tuple[str, str]]:
    """
    Gets the conda environments, as a list of (name, path) tuples.
    """
    prefix = conda_info["conda_prefix"]
    ret = []
    for env in conda_info["envs"]:
        if env == prefix:
            ret.append(("base", env))
        else:
            ret.append((os.path.basename(env), env))
    return ret


class ShellCommandException(Exception):
    pass


def exec_cmd(
    cmd: List[str], throw_on_error: bool = True, logger: Optional[logging.Logger] = None
) -> Union[int, Tuple[int, str, str]]:
    """
    Runs a command as a child process.

    A convenience wrapper for running a command from a Python script.

    Note on the return value: A tuple of the exit code,
    standard output and standard error is returned.

    Args:
        cmd: the command to run, as a list of strings
        throw_on_error: if true, raises an Exception if the exit code of the
            program is nonzero
        logger: Unused; retained for API compatibility.

    Returns:
        A tuple of (exit_code, stdout, stderr) from the child process.
    """
    child = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stdin=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )
    (stdout, stderr) = child.communicate()
    exit_code = child.wait()
    if throw_on_error and exit_code != 0:
        raise ShellCommandException(
            "Non-zero exit code: %s\n\nSTDOUT:\n%s\n\nSTDERR:%s"
            % (exit_code, stdout, stderr)
        )
    return exit_code, stdout, stderr


def exec_cmd_stream_to_logger(
    cmd: List[str], logger: logging.Logger, n_lines: int = 50, **kwargs
) -> Tuple[int, str]:
    """Runs a command as a child process, streaming output to the logger.

    The last n_lines lines of output are also returned (stdout and stderr).
    """
    if "env" in kwargs and _WIN32 and "PATH" not in [x.upper() for x in kwargs.keys]:
        raise ValueError("On windows, Popen requires 'PATH' in 'env'")
    child = subprocess.Popen(
        cmd,
        universal_newlines=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        **kwargs,
    )
    last_n_lines = []
    with child.stdout:
        for line in iter(child.stdout.readline, b""):
            exit_code = child.poll()
            if exit_code is not None:
                break
            line = line.strip()
            if not line:
                continue
            last_n_lines.append(line.strip())
            last_n_lines = last_n_lines[-n_lines:]
            logger.info(line.strip())

    exit_code = child.wait()
    return exit_code, "\n".join(last_n_lines)