import hashlib
import json
import logging
import os

import yaml

from mlflow.exceptions import ExecutionException
from mlflow.utils import process
from mlflow.utils.environment import Environment

# Environment variable indicating a path to a conda installation. MLflow will default to running
# "conda" if unset
MLFLOW_CONDA_HOME = "MLFLOW_CONDA_HOME"
# Environment variable indicated the name of the command that should be used to create environments.
# If it is unset, it will default to "conda". This command must be in the $PATH when the user runs,
# or within MLFLOW_CONDA_HOME if that is set. For example, let's say we want to use mamba
# (https://github.com/mamba-org/mamba) instead of conda to create environments. Then:
# > conda install mamba -n base -c conda-forge
# > MLFLOW_CONDA_CREATE_ENV_CMD="mamba"
# > mlflow run ...
MLFLOW_CONDA_CREATE_ENV_CMD = "MLFLOW_CONDA_CREATE_ENV_CMD"

_logger = logging.getLogger(__name__)


def get_conda_command(conda_env_name):
    #  Checking for newer conda versions
    if os.name != "nt" and ("CONDA_EXE" in os.environ or "MLFLOW_CONDA_HOME" in os.environ):
        conda_path = get_conda_bin_executable("conda")
        activate_conda_env = [
            "source {}/../etc/profile.d/conda.sh".format(os.path.dirname(conda_path))
        ]
        activate_conda_env += [f"conda activate {conda_env_name} 1>&2"]
    else:
        activate_path = get_conda_bin_executable("activate")
        # in case os name is not 'nt', we are not running on windows. It introduces
        # bash command otherwise.
        if os.name != "nt":
            return [f"source {activate_path} {conda_env_name} 1>&2"]
        else:
            return [f"conda activate {conda_env_name}"]
    return activate_conda_env


def get_conda_bin_executable(executable_name):
    """
    Return path to the specified executable, assumed to be discoverable within the 'bin'
    subdirectory of a conda installation.

    The conda home directory (expected to contain a 'bin' subdirectory) is configurable via the
    ``mlflow.projects.MLFLOW_CONDA_HOME`` environment variable. If
    ``mlflow.projects.MLFLOW_CONDA_HOME`` is unspecified, this method simply returns the passed-in
    executable name.
    """
    conda_home = os.environ.get(MLFLOW_CONDA_HOME)
    if conda_home:
        return os.path.join(conda_home, "bin/%s" % executable_name)
    # 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"])
        return os.path.join(conda_bin_dir, executable_name)
    return executable_name


def _get_conda_env_name(conda_env_path, env_id=None, env_root_dir=None):
    if conda_env_path:
        with open(conda_env_path) as f:
            conda_env_contents = f.read()
    else:
        conda_env_contents = ""

    if env_id:
        conda_env_contents += env_id

    env_name = "mlflow-%s" % hashlib.sha1(conda_env_contents.encode("utf-8")).hexdigest()
    if env_root_dir:
        env_root_dir = os.path.normpath(env_root_dir)
        # Generate env name with format "mlflow-{conda_env_contents_hash}-{env_root_dir_hash}"
        # hashing `conda_env_contents` and `env_root_dir` separately helps debugging
        env_name += "-%s" % hashlib.sha1(env_root_dir.encode("utf-8")).hexdigest()

    return env_name


def _get_conda_executable_for_create_env():
    """
    Returns the executable that should be used to create environments. This is "conda"
    by default, but it can be set to something else by setting the environment variable

    """
    conda_env_create_cmd = os.environ.get(MLFLOW_CONDA_CREATE_ENV_CMD)
    if conda_env_create_cmd is not None:
        conda_env_create_path = get_conda_bin_executable(conda_env_create_cmd)
    else:
        # Use the same as conda_path
        conda_env_create_path = get_conda_bin_executable("conda")

    return conda_env_create_path


def _list_conda_environments(extra_env=None):
    """
    Return a list of names of conda environments.

    :param extra_env: extra environment variables for running "conda env list" command.
    """
    prc = process._exec_cmd(
        [get_conda_bin_executable("conda"), "env", "list", "--json"], extra_env=extra_env
    )
    return list(map(os.path.basename, json.loads(prc.stdout).get("envs", [])))


_CONDA_ENVS_DIR = "conda_envs"
_CONDA_CACHE_PKGS_DIR = "conda_cache_pkgs"
_PIP_CACHE_DIR = "pip_cache_pkgs"


def _create_conda_env(
    conda_env_path,
    conda_env_create_path,
    project_env_name,
    conda_extra_env_vars,
    capture_output,
):
    if conda_env_path:
        process._exec_cmd(
            [
                conda_env_create_path,
                "env",
                "create",
                "-n",
                project_env_name,
                "--file",
                conda_env_path,
                "--quiet",
            ],
            extra_env=conda_extra_env_vars,
            capture_output=capture_output,
        )
    else:
        process._exec_cmd(
            [
                conda_env_create_path,
                "create",
                "--channel",
                "conda-forge",
                "--yes",
                "--override-channels",
                "-n",
                project_env_name,
                "python",
            ],
            extra_env=conda_extra_env_vars,
            capture_output=capture_output,
        )

    return Environment(get_conda_command(project_env_name), conda_extra_env_vars)


def _create_conda_env_retry(
    conda_env_path, conda_env_create_path, project_env_name, conda_extra_env_vars, _capture_output
):
    """
    `conda env create` command can fail due to network issues such as `ConnectionResetError`
    while collecting package metadata. This function retries the command up to 3 times.
    """
    num_attempts = 3
    for attempt in range(num_attempts):
        try:
            return _create_conda_env(
                conda_env_path,
                conda_env_create_path,
                project_env_name,
                conda_extra_env_vars,
                capture_output=True,
            )
        except process.ShellCommandException as e:
            error_str = str(e)
            if (num_attempts - attempt - 1) > 0 and (
                "ConnectionResetError" in error_str or "ChunkedEncodingError" in error_str
            ):
                _logger.warning("Conda env creation failed due to network issue. Retrying...")
                continue
            raise


def _get_conda_extra_env_vars(env_root_dir=None):
    """
    Given the `env_root_dir` (See doc of PyFuncBackend constructor argument `env_root_dir`),
    return a dict of environment variables which are used to config conda to generate envs
    under the expected `env_root_dir`.
    """
    if env_root_dir is None:
        return None

    # Create isolated conda package cache dir "conda_pkgs" under the env_root_dir
    # for each python process.
    # Note: shared conda package cache dir causes race condition issues:
    # See https://github.com/conda/conda/issues/8870
    # See https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-environment-directories-envs-dirs
    # and https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-package-directories-pkgs-dirs

    conda_envs_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR)
    conda_pkgs_path = os.path.join(env_root_dir, _CONDA_CACHE_PKGS_DIR)
    pip_cache_dir = os.path.join(env_root_dir, _PIP_CACHE_DIR)

    os.makedirs(conda_envs_path, exist_ok=True)
    os.makedirs(conda_pkgs_path, exist_ok=True)
    os.makedirs(pip_cache_dir, exist_ok=True)

    return {
        "CONDA_ENVS_PATH": conda_envs_path,
        "CONDA_PKGS_DIRS": conda_pkgs_path,
        "PIP_CACHE_DIR": pip_cache_dir,
        # PIP_NO_INPUT=1 makes pip run in non-interactive mode,
        # otherwise pip might prompt "yes or no" and ask stdin input
        "PIP_NO_INPUT": "1",
    }


def get_or_create_conda_env(conda_env_path, env_id=None, capture_output=False, env_root_dir=None):
    """
    Given a `Project`, creates a conda environment containing the project's dependencies if such a
    conda environment doesn't already exist. Returns the name of the conda environment.
    :param conda_env_path: Path to a conda yaml file.
    :param env_id: Optional string that is added to the contents of the yaml file before
                   calculating the hash. It can be used to distinguish environments that have the
                   same conda dependencies but are supposed to be different based on the context.
                   For example, when serving the model we may install additional dependencies to the
                   environment after the environment has been activated.
    :param capture_output: Specify the capture_output argument while executing the
                           "conda env create" command.
    :param env_root_dir: See doc of PyFuncBackend constructor argument `env_root_dir`.
    """

    conda_path = get_conda_bin_executable("conda")
    conda_env_create_path = _get_conda_executable_for_create_env()

    try:
        # Checks if Conda executable exists
        process._exec_cmd([conda_path, "--help"], throw_on_error=False)
    except OSError:
        raise ExecutionException(
            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 MLflow to look for a specific "
            f"Conda executable by setting the {MLFLOW_CONDA_HOME} environment variable "
            "to the path of the Conda executable"
        )

    try:
        # Checks if executable for environment creation exists
        process._exec_cmd([conda_env_create_path, "--help"], throw_on_error=False)
    except OSError:
        raise ExecutionException(
            f"You have set the env variable {MLFLOW_CONDA_CREATE_ENV_CMD}, but "
            f"{conda_env_create_path} does not exist or it is not working properly. "
            f"Note that {conda_env_create_path} and the conda executable need to be "
            "in the same conda environment. You can change the search path by"
            f"modifying the env variable {MLFLOW_CONDA_HOME}"
        )

    conda_extra_env_vars = _get_conda_extra_env_vars(env_root_dir)

    # Include the env_root_dir hash in the project_env_name,
    # this is for avoid conda env name conflicts between different CONDA_ENVS_PATH.
    project_env_name = _get_conda_env_name(conda_env_path, env_id=env_id, env_root_dir=env_root_dir)
    if env_root_dir is not None:
        project_env_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR, project_env_name)
    else:
        project_env_path = project_env_name

    if project_env_name in _list_conda_environments(conda_extra_env_vars):
        _logger.info("Conda environment %s already exists.", project_env_path)
        return Environment(get_conda_command(project_env_name), conda_extra_env_vars)

    _logger.info("=== Creating conda environment %s ===", project_env_path)
    try:
        _create_conda_env_func = (
            # Retry conda env creation in a pytest session to avoid flaky test failures
            _create_conda_env_retry
            if "PYTEST_CURRENT_TEST" in os.environ
            else _create_conda_env
        )
        return _create_conda_env_func(
            conda_env_path,
            conda_env_create_path,
            project_env_name,
            conda_extra_env_vars,
            capture_output,
        )
    except Exception:
        try:
            if project_env_name in _list_conda_environments(conda_extra_env_vars):
                _logger.warning(
                    "Encountered unexpected error while creating conda environment. "
                    "Removing %s.",
                    project_env_path,
                )
                process._exec_cmd(
                    [
                        conda_path,
                        "remove",
                        "--yes",
                        "--name",
                        project_env_name,
                        "--all",
                    ],
                    extra_env=conda_extra_env_vars,
                    capture_output=False,
                )
        except Exception as e:
            _logger.warning(
                "Removing conda environment %s failed (error: %s)",
                project_env_path,
                repr(e),
            )
        raise


def _get_conda_dependencies(conda_yaml_path):
    """
    Extracts conda dependencies from a conda yaml file.

    :param conda_yaml_path: Conda yaml file path.
    """
    with open(conda_yaml_path) as f:
        conda_yaml = yaml.safe_load(f)
        return [d for d in conda_yaml.get("dependencies", []) if isinstance(d, str)]
