import os
import re
from contextlib import contextmanager
from itertools import cycle
from os import PathLike
from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    Iterator,
    List,
    Optional,
    Union,
)
from unittest.mock import Mock

from .config import Config, DataProxy
from .exceptions import Failure, AuthFailure, ResponseNotAccepted
from .runners import Result
from .watchers import FailingResponder

if TYPE_CHECKING:
    from invoke.runners import Runner


class Context(DataProxy):
    """
    Context-aware API wrapper & state-passing object.

    `.Context` objects are created during command-line parsing (or, if desired,
    by hand) and used to share parser and configuration state with executed
    tasks (see :ref:`why-context`).

    Specifically, the class offers wrappers for core API calls (such as `.run`)
    which take into account CLI parser flags, configuration files, and/or
    changes made at runtime. It also acts as a proxy for its `~.Context.config`
    attribute - see that attribute's documentation for details.

    Instances of `.Context` may be shared between tasks when executing
    sub-tasks - either the same context the caller was given, or an altered
    copy thereof (or, theoretically, a brand new one).

    .. versionadded:: 1.0
    """

    def __init__(self, config: Optional[Config] = None) -> None:
        """
        :param config:
            `.Config` object to use as the base configuration.

            Defaults to an anonymous/default `.Config` instance.
        """
        #: The fully merged `.Config` object appropriate for this context.
        #:
        #: `.Config` settings (see their documentation for details) may be
        #: accessed like dictionary keys (``c.config['foo']``) or object
        #: attributes (``c.config.foo``).
        #:
        #: As a convenience shorthand, the `.Context` object proxies to its
        #: ``config`` attribute in the same way - e.g. ``c['foo']`` or
        #: ``c.foo`` returns the same value as ``c.config['foo']``.
        config = config if config is not None else Config()
        self._set(_config=config)
        #: A list of commands to run (via "&&") before the main argument to any
        #: `run` or `sudo` calls. Note that the primary API for manipulating
        #: this list is `prefix`; see its docs for details.
        command_prefixes: List[str] = list()
        self._set(command_prefixes=command_prefixes)
        #: A list of directories to 'cd' into before running commands with
        #: `run` or `sudo`; intended for management via `cd`, please see its
        #: docs for details.
        command_cwds: List[str] = list()
        self._set(command_cwds=command_cwds)

    @property
    def config(self) -> Config:
        # Allows Context to expose a .config attribute even though DataProxy
        # otherwise considers it a config key.
        return self._config

    @config.setter
    def config(self, value: Config) -> None:
        # NOTE: mostly used by client libraries needing to tweak a Context's
        # config at execution time; i.e. a Context subclass that bears its own
        # unique data may want to be stood up when parameterizing/expanding a
        # call list at start of a session, with the final config filled in at
        # runtime.
        self._set(_config=value)

    def run(self, command: str, **kwargs: Any) -> Optional[Result]:
        """
        Execute a local shell command, honoring config options.

        Specifically, this method instantiates a `.Runner` subclass (according
        to the ``runner`` config option; default is `.Local`) and calls its
        ``.run`` method with ``command`` and ``kwargs``.

        See `.Runner.run` for details on ``command`` and the available keyword
        arguments.

        .. versionadded:: 1.0
        """
        runner = self.config.runners.local(self)
        return self._run(runner, command, **kwargs)

    # NOTE: broken out of run() to allow for runner class injection in
    # Fabric/etc, which needs to juggle multiple runner class types (local and
    # remote).
    def _run(
        self, runner: "Runner", command: str, **kwargs: Any
    ) -> Optional[Result]:
        command = self._prefix_commands(command)
        return runner.run(command, **kwargs)

    def sudo(self, command: str, **kwargs: Any) -> Optional[Result]:
        """
        Execute a shell command via ``sudo`` with password auto-response.

        **Basics**

        This method is identical to `run` but adds a handful of
        convenient behaviors around invoking the ``sudo`` program. It doesn't
        do anything users could not do themselves by wrapping `run`, but the
        use case is too common to make users reinvent these wheels themselves.

        .. note::
            If you intend to respond to sudo's password prompt by hand, just
            use ``run("sudo command")`` instead! The autoresponding features in
            this method will just get in your way.

        Specifically, `sudo`:

        * Places a `.FailingResponder` into the ``watchers`` kwarg (see
          :doc:`/concepts/watchers`) which:

            * searches for the configured ``sudo`` password prompt;
            * responds with the configured sudo password (``sudo.password``
              from the :doc:`configuration </concepts/configuration>`);
            * can tell when that response causes an authentication failure
              (e.g. if the system requires a password and one was not
              configured), and raises `.AuthFailure` if so.

        * Builds a ``sudo`` command string using the supplied ``command``
          argument, prefixed by various flags (see below);
        * Executes that command via a call to `run`, returning the result.

        **Flags used**

        ``sudo`` flags used under the hood include:

        - ``-S`` to allow auto-responding of password via stdin;
        - ``-p <prompt>`` to explicitly state the prompt to use, so we can be
          sure our auto-responder knows what to look for;
        - ``-u <user>`` if ``user`` is not ``None``, to execute the command as
          a user other than ``root``;
        - When ``-u`` is present, ``-H`` is also added, to ensure the
          subprocess has the requested user's ``$HOME`` set properly.

        **Configuring behavior**

        There are a couple of ways to change how this method behaves:

        - Because it wraps `run`, it honors all `run` config parameters and
          keyword arguments, in the same way that `run` does.

            - Thus, invocations such as ``c.sudo('command', echo=True)`` are
              possible, and if a config layer (such as a config file or env
              var) specifies that e.g. ``run.warn = True``, that too will take
              effect under `sudo`.

        - `sudo` has its own set of keyword arguments (see below) and they are
          also all controllable via the configuration system, under the
          ``sudo.*`` tree.

            - Thus you could, for example, pre-set a sudo user in a config
              file; such as an ``invoke.json`` containing ``{"sudo": {"user":
              "someuser"}}``.

        :param str password: Runtime override for ``sudo.password``.
        :param str user: Runtime override for ``sudo.user``.

        .. versionadded:: 1.0
        """
        runner = self.config.runners.local(self)
        return self._sudo(runner, command, **kwargs)

    # NOTE: this is for runner injection; see NOTE above _run().
    def _sudo(
        self, runner: "Runner", command: str, **kwargs: Any
    ) -> Optional[Result]:
        prompt = self.config.sudo.prompt
        password = kwargs.pop("password", self.config.sudo.password)
        user = kwargs.pop("user", self.config.sudo.user)
        env = kwargs.get("env", {})
        # TODO: allow subclassing for 'get the password' so users who REALLY
        # want lazy runtime prompting can have it easily implemented.
        # TODO: want to print a "cleaner" echo with just 'sudo <command>'; but
        # hard to do as-is, obtaining config data from outside a Runner one
        # holds is currently messy (could fix that), if instead we manually
        # inspect the config ourselves that duplicates logic. NOTE: once we
        # figure that out, there is an existing, would-fail-if-not-skipped test
        # for this behavior in test/context.py.
        # TODO: once that is done, though: how to handle "full debug" output
        # exactly (display of actual, real full sudo command w/ -S and -p), in
        # terms of API/config? Impl is easy, just go back to passing echo
        # through to 'run'...
        user_flags = ""
        if user is not None:
            user_flags = "-H -u {} ".format(user)
        env_flags = ""
        if env:
            env_flags = "--preserve-env='{}' ".format(",".join(env.keys()))
        command = self._prefix_commands(command)
        cmd_str = "sudo -S -p '{}' {}{}{}".format(
            prompt, env_flags, user_flags, command
        )
        watcher = FailingResponder(
            pattern=re.escape(prompt),
            response="{}\n".format(password),
            sentinel="Sorry, try again.\n",
        )
        # Ensure we merge any user-specified watchers with our own.
        # NOTE: If there are config-driven watchers, we pull those up to the
        # kwarg level; that lets us merge cleanly without needing complex
        # config-driven "override vs merge" semantics.
        # TODO: if/when those semantics are implemented, use them instead.
        # NOTE: config value for watchers defaults to an empty list; and we
        # want to clone it to avoid actually mutating the config.
        watchers = kwargs.pop("watchers", list(self.config.run.watchers))
        watchers.append(watcher)
        try:
            return runner.run(cmd_str, watchers=watchers, **kwargs)
        except Failure as failure:
            # Transmute failures driven by our FailingResponder, into auth
            # failures - the command never even ran.
            # TODO: wants to be a hook here for users that desire "override a
            # bad config value for sudo.password" manual input
            # NOTE: as noted in #294 comments, we MAY in future want to update
            # this so run() is given ability to raise AuthFailure on its own.
            # For now that has been judged unnecessary complexity.
            if isinstance(failure.reason, ResponseNotAccepted):
                # NOTE: not bothering with 'reason' here, it's pointless.
                error = AuthFailure(result=failure.result, prompt=prompt)
                raise error
            # Reraise for any other error so it bubbles up normally.
            else:
                raise

    # TODO: wonder if it makes sense to move this part of things inside Runner,
    # which would grow a `prefixes` and `cwd` init kwargs or similar. The less
    # that's stuffed into Context, probably the better.
    def _prefix_commands(self, command: str) -> str:
        """
        Prefixes ``command`` with all prefixes found in ``command_prefixes``.

        ``command_prefixes`` is a list of strings which is modified by the
        `prefix` context manager.
        """
        prefixes = list(self.command_prefixes)
        current_directory = self.cwd
        if current_directory:
            prefixes.insert(0, "cd {}".format(current_directory))

        return " && ".join(prefixes + [command])

    @contextmanager
    def prefix(self, command: str) -> Generator[None, None, None]:
        """
        Prefix all nested `run`/`sudo` commands with given command plus ``&&``.

        Most of the time, you'll want to be using this alongside a shell script
        which alters shell state, such as ones which export or alter shell
        environment variables.

        For example, one of the most common uses of this tool is with the
        ``workon`` command from `virtualenvwrapper
        <https://virtualenvwrapper.readthedocs.io/en/latest/>`_::

            with c.prefix('workon myvenv'):
                c.run('./manage.py migrate')

        In the above snippet, the actual shell command run would be this::

            $ workon myvenv && ./manage.py migrate

        This context manager is compatible with `cd`, so if your virtualenv
        doesn't ``cd`` in its ``postactivate`` script, you could do the
        following::

            with c.cd('/path/to/app'):
                with c.prefix('workon myvenv'):
                    c.run('./manage.py migrate')
                    c.run('./manage.py loaddata fixture')

        Which would result in executions like so::

            $ cd /path/to/app && workon myvenv && ./manage.py migrate
            $ cd /path/to/app && workon myvenv && ./manage.py loaddata fixture

        Finally, as alluded to above, `prefix` may be nested if desired, e.g.::

            with c.prefix('workon myenv'):
                c.run('ls')
                with c.prefix('source /some/script'):
                    c.run('touch a_file')

        The result::

            $ workon myenv && ls
            $ workon myenv && source /some/script && touch a_file

        Contrived, but hopefully illustrative.

        .. versionadded:: 1.0
        """
        self.command_prefixes.append(command)
        try:
            yield
        finally:
            self.command_prefixes.pop()

    @property
    def cwd(self) -> str:
        """
        Return the current working directory, accounting for uses of `cd`.

        .. versionadded:: 1.0
        """
        if not self.command_cwds:
            # TODO: should this be None? Feels cleaner, though there may be
            # benefits to it being an empty string, such as relying on a no-arg
            # `cd` typically being shorthand for "go to user's $HOME".
            return ""

        # get the index for the subset of paths starting with the last / or ~
        for i, path in reversed(list(enumerate(self.command_cwds))):
            if path.startswith("~") or path.startswith("/"):
                break

        # TODO: see if there's a stronger "escape this path" function somewhere
        # we can reuse. e.g., escaping tildes or slashes in filenames.
        paths = [path.replace(" ", r"\ ") for path in self.command_cwds[i:]]
        return str(os.path.join(*paths))

    @contextmanager
    def cd(self, path: Union[PathLike, str]) -> Generator[None, None, None]:
        """
        Context manager that keeps directory state when executing commands.

        Any calls to `run`, `sudo`, within the wrapped block will implicitly
        have a string similar to ``"cd <path> && "`` prefixed in order to give
        the sense that there is actually statefulness involved.

        Because use of `cd` affects all such invocations, any code making use
        of the `cwd` property will also be affected by use of `cd`.

        Like the actual 'cd' shell builtin, `cd` may be called with relative
        paths (keep in mind that your default starting directory is your user's
        ``$HOME``) and may be nested as well.

        Below is a "normal" attempt at using the shell 'cd', which doesn't work
        since all commands are executed in individual subprocesses -- state is
        **not** kept between invocations of `run` or `sudo`::

            c.run('cd /var/www')
            c.run('ls')

        The above snippet will list the contents of the user's ``$HOME``
        instead of ``/var/www``. With `cd`, however, it will work as expected::

            with c.cd('/var/www'):
                c.run('ls')  # Turns into "cd /var/www && ls"

        Finally, a demonstration (see inline comments) of nesting::

            with c.cd('/var/www'):
                c.run('ls') # cd /var/www && ls
                with c.cd('website1'):
                    c.run('ls')  # cd /var/www/website1 && ls

        .. note::
            Space characters will be escaped automatically to make dealing with
            such directory names easier.

        .. versionadded:: 1.0
        .. versionchanged:: 1.5
            Explicitly cast the ``path`` argument (the only argument) to a
            string; this allows any object defining ``__str__`` to be handed in
            (such as the various ``Path`` objects out there), and not just
            string literals.
        """
        path = str(path)
        self.command_cwds.append(path)
        try:
            yield
        finally:
            self.command_cwds.pop()


class MockContext(Context):
    """
    A `.Context` whose methods' return values can be predetermined.

    Primarily useful for testing Invoke-using codebases.

    .. note::
        This class wraps its ``run``, etc methods in `unittest.mock.Mock`
        objects. This allows you to easily assert that the methods (still
        returning the values you prepare them with) were actually called.

    .. note::
        Methods not given `Results <.Result>` to yield will raise
        ``NotImplementedError`` if called (since the alternative is to call the
        real underlying method - typically undesirable when mocking.)

    .. versionadded:: 1.0
    .. versionchanged:: 1.5
        Added ``Mock`` wrapping of ``run`` and ``sudo``.
    """

    def __init__(self, config: Optional[Config] = None, **kwargs: Any) -> None:
        """
        Create a ``Context``-like object whose methods yield `.Result` objects.

        :param config:
            A Configuration object to use. Identical in behavior to `.Context`.

        :param run:
            A data structure indicating what `.Result` objects to return from
            calls to the instantiated object's `~.Context.run` method (instead
            of actually executing the requested shell command).

            Specifically, this kwarg accepts:

            - A single `.Result` object.
            - A boolean; if True, yields a `.Result` whose ``exited`` is ``0``,
              and if False, ``1``.
            - An iterable of the above values, which will be returned on each
              subsequent call to ``.run`` (the first item on the first call,
              the second on the second call, etc).
            - A dict mapping command strings or compiled regexen to the above
              values (including an iterable), allowing specific
              call-and-response semantics instead of assuming a call order.

        :param sudo:
            Identical to ``run``, but whose values are yielded from calls to
            `~.Context.sudo`.

        :param bool repeat:
            A flag determining whether results yielded by this class' methods
            repeat or are consumed.

            For example, when a single result is indicated, it will normally
            only be returned once, causing ``NotImplementedError`` afterwards.
            But when ``repeat=True`` is given, that result is returned on
            every call, forever.

            Similarly, iterable results are normally exhausted once, but when
            this setting is enabled, they are wrapped in `itertools.cycle`.

            Default: ``True``.

        :raises:
            ``TypeError``, if the values given to ``run`` or other kwargs
            aren't of the expected types.

        .. versionchanged:: 1.5
            Added support for boolean and string result values.
        .. versionchanged:: 1.5
            Added support for regex dict keys.
        .. versionchanged:: 1.5
            Added the ``repeat`` keyword argument.
        .. versionchanged:: 2.0
            Changed ``repeat`` default value from ``False`` to ``True``.
        """
        # Set up like any other Context would, with the config
        super().__init__(config)
        # Pull out behavioral kwargs
        self._set("__repeat", kwargs.pop("repeat", True))
        # The rest must be things like run/sudo - mock Context method info
        for method, results in kwargs.items():
            # For each possible value type, normalize to iterable of Result
            # objects (possibly repeating).
            singletons = (Result, bool, str)
            if isinstance(results, dict):
                for key, value in results.items():
                    results[key] = self._normalize(value)
            elif isinstance(results, singletons) or hasattr(
                results, "__iter__"
            ):
                results = self._normalize(results)
            # Unknown input value: cry
            else:
                err = "Not sure how to yield results from a {!r}"
                raise TypeError(err.format(type(results)))
            # Save results for use by the method
            self._set("__{}".format(method), results)
            # Wrap the method in a Mock
            self._set(method, Mock(wraps=getattr(self, method)))

    def _normalize(self, value: Any) -> Iterator[Any]:
        # First turn everything into an iterable
        if not hasattr(value, "__iter__") or isinstance(value, str):
            value = [value]
        # Then turn everything within into a Result
        results = []
        for obj in value:
            if isinstance(obj, bool):
                obj = Result(exited=0 if obj else 1)
            elif isinstance(obj, str):
                obj = Result(obj)
            results.append(obj)
        # Finally, turn that iterable into an iteratOR, depending on repeat
        return cycle(results) if getattr(self, "__repeat") else iter(results)

    # TODO: _maybe_ make this more metaprogrammy/flexible (using __call__ etc)?
    # Pretty worried it'd cause more hard-to-debug issues than it's presently
    # worth. Maybe in situations where Context grows a _lot_ of methods (e.g.
    # in Fabric 2; though Fabric could do its own sub-subclass in that case...)

    def _yield_result(self, attname: str, command: str) -> Result:
        try:
            obj = getattr(self, attname)
            # Dicts need to try direct lookup or regex matching
            if isinstance(obj, dict):
                try:
                    obj = obj[command]
                except KeyError:
                    # TODO: could optimize by skipping this if not any regex
                    # objects in keys()?
                    for key, value in obj.items():
                        if hasattr(key, "match") and key.match(command):
                            obj = value
                            break
                    else:
                        # Nope, nothing did match.
                        raise KeyError
            # Here, the value was either never a dict or has been extracted
            # from one, so we can assume it's an iterable of Result objects due
            # to work done by __init__.
            result: Result = next(obj)
            # Populate Result's command string with what matched unless
            # explicitly given
            if not result.command:
                result.command = command
            return result
        except (AttributeError, IndexError, KeyError, StopIteration):
            # raise_from(NotImplementedError(command), None)
            raise NotImplementedError(command)

    def run(self, command: str, *args: Any, **kwargs: Any) -> Result:
        # TODO: perform more convenience stuff associating args/kwargs with the
        # result? E.g. filling in .command, etc? Possibly useful for debugging
        # if one hits unexpected-order problems with what they passed in to
        # __init__.
        return self._yield_result("__run", command)

    def sudo(self, command: str, *args: Any, **kwargs: Any) -> Result:
        # TODO: this completely nukes the top-level behavior of sudo(), which
        # could be good or bad, depending. Most of the time I think it's good.
        # No need to supply dummy password config, etc.
        # TODO: see the TODO from run() re: injecting arg/kwarg values
        return self._yield_result("__sudo", command)

    def set_result_for(
        self, attname: str, command: str, result: Result
    ) -> None:
        """
        Modify the stored mock results for given ``attname`` (e.g. ``run``).

        This is similar to how one instantiates `MockContext` with a ``run`` or
        ``sudo`` dict kwarg. For example, this::

            mc = MockContext(run={'mycommand': Result("mystdout")})
            assert mc.run('mycommand').stdout == "mystdout"

        is functionally equivalent to this::

            mc = MockContext()
            mc.set_result_for('run', 'mycommand', Result("mystdout"))
            assert mc.run('mycommand').stdout == "mystdout"

        `set_result_for` is mostly useful for modifying an already-instantiated
        `MockContext`, such as one created by test setup or helper methods.

        .. versionadded:: 1.0
        """
        attname = "__{}".format(attname)
        heck = TypeError(
            "Can't update results for non-dict or nonexistent mock results!"
        )
        # Get value & complain if it's not a dict.
        # TODO: should we allow this to set non-dict values too? Seems vaguely
        # pointless, at that point, just make a new MockContext eh?
        try:
            value = getattr(self, attname)
        except AttributeError:
            raise heck
        if not isinstance(value, dict):
            raise heck
        # OK, we're good to modify, so do so.
        value[command] = self._normalize(result)