# Licensed under a 3-clause BSD style license - see LICENSE.rst
import contextlib
import os
import re
import tempfile
from pathlib import Path

from packaging.version import Version

from .. import environment, util
from ..console import log

WIN = os.name == "nt"


# Conda (as of version 4.7.5) is not safe to run in parallel.
# See https://github.com/conda/conda/issues/8870
# Hence, serialize the calls to it.

util.new_multiprocessing_lock("conda_lock")


def _conda_lock():
    # function; for easier monkeypatching
    return util.get_multiprocessing_lock("conda_lock")


@contextlib.contextmanager
def _dummy_lock():
    yield


def _find_conda():
    """
    Find the conda executable robustly across conda versions.

    Returns
    -------
    conda : str
        Path to the conda executable.

    Raises
    ------
    OSError
        If the executable cannot be found in either the CONDA_EXE environment
        variable or in the PATH.

    Notes
    -----
    In POSIX platforms in conda >= 4.4, conda can be set up as a bash function
    rather than an executable. (This is to enable the syntax
    ``conda activate env-name``.) In this case, the environment variable
    ``CONDA_EXE`` contains the path to the conda executable. In other cases,
    we use standard search for the appropriate name in the PATH.

    See https://github.com/airspeed-velocity/asv/issues/645 for more details.
    """
    if 'CONDA_EXE' in os.environ:
        conda = os.environ['CONDA_EXE']
    else:
        conda = util.which('conda')
    return conda


class Conda(environment.Environment):
    """
    Manage an environment using conda.

    Dependencies are installed using ``conda``.  The benchmarked
    project is installed using ``pip``.
    """

    tool_name = "conda"
    _matches_cache = {}

    def __init__(self, conf, python, requirements, tagged_env_vars):
        """
        Parameters
        ----------
        conf : Config instance

        python : str
            Version of Python.  Must be of the form "MAJOR.MINOR".

        requirements : dict
            Dictionary mapping a PyPI package name to a version
            identifier string.
        """
        self._python = python
        self._requirements = requirements
        self._conda_channels = conf.conda_channels
        if "conda-forge" not in conf.conda_channels:
            self._conda_channels += ["conda-forge"]
        self._conda_environment_file = conf.conda_environment_file

        if conf.conda_environment_file == "IGNORE":
            log.debug("Skipping environment file due to conda_environment_file set to IGNORE")
            self._conda_environment_file = None
        elif not conf.conda_environment_file:
            if Path("environment.yml").exists():
                log.debug("Using environment.yml")
                self._conda_environment_file = "environment.yml"

        super().__init__(conf, python, requirements, tagged_env_vars)

    @classmethod
    def matches(cls, python):
        # Calling conda can take a long time, so remember the result
        if python not in cls._matches_cache:
            cls._matches_cache[python] = cls._matches(python)
        return cls._matches_cache[python]

    @classmethod
    def _matches(cls, python):
        if not re.match(r'^[0-9].*$', python):
            return False
        else:
            conda = _find_conda()
            try:
                with _conda_lock():
                    return util.search_channels(conda, "python", python)
            except util.ProcessError:
                return False

    def _setup(self):
        log.info(f"Creating conda environment for {self.name}")

        conda_args, pip_args = self._get_requirements()
        env = dict(os.environ)
        env.update(self.build_env_vars)

        # Changed in v0.6.5, gh-1294
        # previously, the user provided environment was assumed to handle the python version
        conda_args = [util.replace_cpython_version(arg, self._python) for arg in conda_args]

        if not self._conda_environment_file:
            # With a user-provided envronment, we assume it specifies a python version;
            # without an environment.yml file, we need to add the python version ourselves
            conda_args = [f'python={self._python}', 'wheel', 'pip'] + conda_args

        # Create a temporary environment.yml file
        # and use that to generate the env for benchmarking.
        env_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".yml")
        try:
            env_file.write(f'name: {self.name}\nchannels:\n')
            env_file.writelines(f'   - {ch}\n' for ch in self._conda_channels)
            if conda_args:
                env_file.write('dependencies:\n')
                # categorize & write dependencies based on pip vs. conda
                env_file.writelines(f'   - {s}\n' for s in conda_args)
            env_file.close()
            try:
                env_file_name = self._conda_environment_file or env_file.name

                conda_version = re.search(r'\d+(\.\d+)+', self._run_conda(['--version'], env=env))[
                    0
                ]
                log.info(f"conda version: {conda_version}")
                # https://conda.io/projects/conda/en/latest/release-notes.html#id8
                if Version(conda_version) >= Version("24.3.0"):
                    self._run_conda(
                        ['env', 'create', '-f', env_file_name, '-p', self._path, "--yes"], env=env
                    )
                else:  # Backward compatibility
                    self._run_conda(
                        ['env', 'create', '-f', env_file_name, '-p', self._path, '--force'],
                        env=env,
                    )

                if self._conda_environment_file and (conda_args or pip_args):
                    # Add extra packages
                    env_file_name = env_file.name
                    self._run_conda(
                        ['env', 'update', '-f', env_file_name, '-p', self._path], env=env
                    )
            except Exception:
                if env_file_name != env_file.name:
                    log.info(
                        "conda env create/update failed: "
                        f"in {self._path} with file {env_file_name}"
                    )
                elif os.path.isfile(env_file_name):
                    with open(env_file_name, 'r') as f:
                        text = f.read()
                    log.info(f"conda env create/update failed: in {self._path} with:\n{text}")
                raise
        finally:
            os.unlink(env_file.name)
        if pip_args:
            for declaration in pip_args:
                parsed_declaration = util.ParsedPipDeclaration(declaration)
                pip_call = util.construct_pip_call(self._run_pip, parsed_declaration)
                pip_call()

    def _get_requirements(self):
        conda_args = []
        pip_args = []

        for key, val in {**self._requirements, **self._base_requirements}.items():
            if key.startswith("pip+"):
                pip_args.append(f"{key[4:]} {val}")
            else:
                if val:
                    conda_args.append(f"{key}={val}")
                else:
                    conda_args.append(key)

        return conda_args, pip_args

    def _run_conda(self, args, env=None):
        """
        Run conda command outside the environment.
        """
        try:
            conda = _find_conda()
        except OSError as e:
            raise util.UserError(str(e))

        with _conda_lock():
            return util.check_output([conda] + args, timeout=self._install_timeout, env=env)

    def run(self, args, **kwargs):
        log.debug(f"Running '{' '.join(args)}' in {self.name}")
        return self.run_executable('python', args, **kwargs)

    def run_executable(self, executable, args, **kwargs):
        # Special-case running conda, for user-provided commands
        if executable == "conda":
            executable = _find_conda()
            lock = _conda_lock
        else:
            lock = _dummy_lock

        # Conda doesn't guarantee that user site directories are excluded
        kwargs["env"] = dict(kwargs.pop("env", os.environ), PYTHONNOUSERSITE="True")

        with lock():
            return super().run_executable(executable, args, **kwargs)

    def _run_pip(self, args, **kwargs):
        # Run pip via python -m pip, so that it works on Windows when
        # upgrading pip itself, and avoids shebang length limit on Linux
        return self.run_executable("python", ["-m", "pip"] + list(args), **kwargs)
