# Copyright 2020 Pex project contributors.
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import, print_function

import functools
import json
import logging
import os
import subprocess
import sys
import tempfile
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace, _ActionsContainer
from contextlib import contextmanager
from subprocess import CalledProcessError

from pex import pex_warnings
from pex.argparse import HandleBoolAction
from pex.cache import access as cache_access
from pex.common import environment_as, safe_mkdtemp, safe_open
from pex.compatibility import shlex_quote
from pex.os import MAC, WINDOWS
from pex.result import Error, Ok, Result
from pex.subprocess import subprocess_daemon_kwargs
from pex.typing import TYPE_CHECKING, Generic, cast, overload
from pex.variables import ENV, Variables
from pex.version import __version__

if TYPE_CHECKING:
    from typing import (
        IO,
        Any,
        ContextManager,
        Dict,
        Iterable,
        Iterator,
        NoReturn,
        Optional,
        Sequence,
        Tuple,
        Type,
        TypeVar,
    )

    import attr  # vendor:skip
    from typing_extensions import Literal
else:
    from pex.third_party import attr

if TYPE_CHECKING:
    _T = TypeVar("_T")


def try_run_program(
    program,  # type: str
    args,  # type: Iterable[str]
    program_info_url=None,  # type: Optional[str]
    error=None,  # type: Optional[str]
    disown=False,  # type: bool
    **kwargs  # type: Any
):
    # type: (...) -> Result
    cmd = [program] + list(args)
    kwargs = dict(subprocess_daemon_kwargs() if disown else {}, **kwargs)
    try:
        process = subprocess.Popen(cmd, **kwargs)
        if not disown and process.wait() != 0:
            return Error(
                str(
                    CalledProcessError(
                        returncode=process.returncode,
                        cmd=" ".join(shlex_quote(arg) for arg in cmd),
                    )
                ),
                exit_code=process.returncode,
            )
        return Ok()
    except OSError as e:
        msg = [error] if error else []
        msg.append("Do you have `{}` installed on the $PATH?: {}".format(program, e))
        if program_info_url:
            msg.append(
                "Find more information on `{program}` at {url}.".format(
                    program=program, url=program_info_url
                )
            )
        return Error("\n".join(msg))


def try_open(
    path_or_url,  # type: str
    open_program=None,  # type: Optional[str]
    error=None,  # type: Optional[str]
    suppress_stderr=False,  # type: bool
):
    # type: (...) -> Result

    program_info_url = None  # type: Optional[str]
    if open_program:
        opener = open_program
    elif WINDOWS:
        opener = "explorer"
    elif MAC:
        opener = "open"
    else:
        opener = "xdg-open"
        program_info_url = "https://www.freedesktop.org/wiki/Software/xdg-utils/"

    with open(os.devnull, "wb") as devnull:
        return try_run_program(
            opener,
            [path_or_url],
            program_info_url=program_info_url,
            error=error,
            disown=True,
            stdout=devnull,
            stderr=devnull if suppress_stderr else None,
        )


@attr.s(frozen=True)
class Command(object):
    @staticmethod
    def show_help(
        parser,  # type: ArgumentParser
        *_args,  # type: Any
        **_kwargs  # type: Any
    ):
        # type: (...) -> NoReturn
        parser.error("a subcommand is required")

    @staticmethod
    def register_global_arguments(
        parser,  # type: _ActionsContainer
        include_verbosity=True,  # type: bool
    ):
        # type: (...) -> None
        register_global_arguments(parser, include_verbosity=include_verbosity)

    @classmethod
    def name(cls):
        # type: () -> str
        return cls.__name__.lower()

    @classmethod
    def description(cls):
        # type: () -> Optional[str]
        return cls.__doc__

    @classmethod
    def add_arguments(cls, parser):
        # type: (ArgumentParser) -> None
        pass

    @classmethod
    def supports_unknown_args(cls):
        # type: () -> bool
        return False

    options = attr.ib()  # type: Namespace
    passthrough_args = attr.ib(default=None)  # type: Optional[Tuple[str, ...]]


class OutputMixin(object):
    @staticmethod
    def add_output_option(
        parser,  # type: _ActionsContainer
        entity,  # type: str
    ):
        # type: (...) -> None
        parser.add_argument(
            "-o",
            "--output",
            metavar="PATH",
            help=(
                "A file to output the {entity} to; STDOUT by default or when `-` is "
                "specified.".format(entity=entity)
            ),
        )

    @staticmethod
    def is_stdout(options):
        # type: (Namespace) -> bool
        return options.output == "-" or not options.output

    @overload
    @classmethod
    @contextmanager
    def output(
        cls,
        options,  # type: Namespace
    ):
        # type: (...) -> ContextManager[IO[str]]
        pass

    @overload
    @classmethod
    @contextmanager
    def output(
        cls,
        options,  # type: Namespace
        binary,  # type: Literal[False]
    ):
        # type: (...) -> ContextManager[IO[str]]
        pass

    @overload
    @classmethod
    @contextmanager
    def output(
        cls,
        options,  # type: Namespace
        binary,  # type: Literal[True]
    ):
        # type: (...) -> ContextManager[IO[bytes]]
        pass

    @classmethod
    @contextmanager
    def output(
        cls,
        options,  # type: Namespace
        binary=False,  # type: bool
    ):
        # type: (...) -> Iterator[IO]
        if cls.is_stdout(options):
            stdout = getattr(sys.stdout, "buffer", sys.stdout) if binary else sys.stdout
            yield stdout
        else:
            with safe_open(options.output, mode="wb" if binary else "w") as out:
                yield out


class JsonMixin(object):
    @staticmethod
    def add_json_options(
        parser,  # type: _ActionsContainer
        entity,  # type: str
        include_switch=True,  # type: bool
    ):
        flags = ("-i", "--indent") if include_switch else ("--indent",)
        parser.add_argument(
            *flags,
            type=int,
            default=None,
            help="Pretty-print {entity} json with the given indent.".format(entity=entity)
        )

    @staticmethod
    def dump_json(
        options,  # type: Namespace
        data,  # type: Dict[str, Any]
        out,  # type: IO
        **json_dump_kwargs  # type: Any
    ):
        if options.indent is not None and options.indent > 0:
            # Python 2.7 uses ', ' for the list item separator regardless of indent which is
            # different from Python 3 and leads to trailing whitespace in the output; so, we
            # normalize here to the Python 3 style for consistent, more generally useful output.
            json_dump_kwargs.update(separators=(",", ": "))
        json.dump(data, out, indent=options.indent, **json_dump_kwargs)


def register_global_arguments(
    parser,  # type: _ActionsContainer
    include_verbosity=True,  # type: bool
):
    # type: (...) -> None
    """Register Pex global environment configuration options with the given parser.

    :param parser: The parser to register global options with.
    :param include_verbosity: Whether to include the verbosity option `-v`.
    """

    group = parser.add_argument_group(title="Global options")
    if include_verbosity:
        group.add_argument(
            "-v",
            dest="verbosity",
            action="count",
            default=0,
            help="Turn on logging verbosity, may be specified multiple times.",
        )
    group.add_argument(
        "--emit-warnings",
        "--no-emit-warnings",
        dest="emit_warnings",
        action=HandleBoolAction,
        default=True,
        help=(
            "Emit runtime UserWarnings on stderr. If false, only emit them when PEX_VERBOSE "
            "is set."
        ),
    )
    group.add_argument(
        "--pex-root",
        dest="pex_root",
        default=None,
        help=(
            "Specify the pex root used in this invocation of pex "
            "(if unspecified, uses {}).".format(ENV.PEX_ROOT)
        ),
    )
    group.add_argument(
        "--disable-cache",
        dest="disable_cache",
        default=False,
        action="store_true",
        help="Disable caching in the pex tool entirely.",
    )

    group.add_argument(
        "--cache-dir",
        dest="cache_dir",
        default=None,
        help=(
            "DEPRECATED: Use --pex-root instead. The local cache directory to use for speeding up "
            "requirement lookups."
        ),
    )
    group.add_argument(
        "--tmpdir",
        dest="tmpdir",
        default=tempfile.gettempdir(),
        help="Specify the temporary directory Pex and its subprocesses should use.",
    )
    group.add_argument(
        "--rcfile",
        dest="rc_file",
        default=None,
        help=(
            "An additional path to a pexrc file to read during configuration parsing, in addition "
            "to reading `/etc/pexrc` and `~/.pexrc`. If `PEX_IGNORE_RCFILES=true`, then all rc "
            "files will be ignored."
        ),
    )


class GlobalConfigurationError(Exception):
    """Indicates an error processing global options."""


@contextmanager
def _configured_env(options):
    # type: (Namespace) -> Iterator[None]
    if options.rc_file or not ENV.PEX_IGNORE_RCFILES:
        with ENV.patch(**Variables(rc=options.rc_file).copy()):
            yield
    else:
        yield


@contextmanager
def global_environment(options):
    # type: (Namespace) -> Iterator[Dict[str, str]]
    """Configures the Pex global environment.

    This includes configuration of basic Pex infrastructure like logging, warnings and the
    `PEX_ROOT` to use.

    :param options: The global options registered by `register_global_arguments`.
    :yields: The configured global environment.
    :raises: :class:`GlobalConfigurationError` if invalid global option values were specified.
    """
    if not hasattr(options, "rc_file"):
        # We don't register the global args on the root command (but do on every subcommand).
        # So if the user runs just `pex` with no subcommand we must not attempt to use those
        # global args, including rc_file, which we check for here as a representative of the
        # global args.
        # Note that we can't use command_type here because the legacy command line parser in
        # pex/bin/pex.py uses this function as well, and it doesn't set command_type.
        with ENV.patch() as env:
            yield env
    with _configured_env(options):
        verbosity = Variables.PEX_VERBOSE.strip_default(ENV)
        if verbosity is None:
            verbosity = getattr(options, "verbosity", 0)

        emit_warnings = True
        if not options.emit_warnings:
            emit_warnings = False
        if emit_warnings and ENV.PEX_EMIT_WARNINGS is not None:
            emit_warnings = ENV.PEX_EMIT_WARNINGS

        with ENV.patch(PEX_VERBOSE=str(verbosity), PEX_EMIT_WARNINGS=str(emit_warnings)):
            pex_warnings.configure_warnings(env=ENV)

            # Ensure the TMPDIR is an absolute path (So subprocesses that change CWD can find it)
            # and that it exists.
            tmpdir = os.path.realpath(options.tmpdir)
            if not os.path.exists(tmpdir):
                raise GlobalConfigurationError(
                    "The specified --tmpdir does not exist: {}".format(tmpdir)
                )
            if not os.path.isdir(tmpdir):
                raise GlobalConfigurationError(
                    "The specified --tmpdir is not a directory: {}".format(tmpdir)
                )
            tempfile.tempdir = os.environ["TMPDIR"] = tmpdir

            if options.cache_dir:
                pex_warnings.warn("The --cache-dir option is deprecated, use --pex-root instead.")
                if options.pex_root and options.cache_dir != options.pex_root:
                    raise GlobalConfigurationError(
                        "Both --cache-dir and --pex-root were passed with conflicting values. "
                        "Just set --pex-root."
                    )

            if options.disable_cache:

                def warn_ignore_pex_root(set_via):
                    pex_warnings.warn(
                        "The pex root has been set via {via} but --disable-cache is also set. "
                        "Ignoring {via} and disabling caches.".format(via=set_via)
                    )

                if options.cache_dir:
                    warn_ignore_pex_root("--cache-dir")
                elif options.pex_root:
                    warn_ignore_pex_root("--pex-root")
                elif os.environ.get("PEX_ROOT"):
                    warn_ignore_pex_root("PEX_ROOT")

                pex_root = safe_mkdtemp()
            else:
                pex_root = options.cache_dir or options.pex_root or ENV.PEX_ROOT

            with ENV.patch(PEX_ROOT=pex_root, TMPDIR=tmpdir) as env, environment_as(**env):
                cache_access.read_write()
                yield env


if TYPE_CHECKING:
    _C = TypeVar("_C", bound=Command)


class Main(Generic["_C"]):
    def __init__(
        self,
        command_types,  # type: Iterable[Type[_C]]
        description=None,  # type: Optional[str]
        subparsers_description=None,  # type: Optional[str]
        prog=None,  # type: Optional[str]
    ):
        # type: (...) -> None
        self._prog = prog
        self._description = description or self.__doc__
        self._subparsers_description = subparsers_description
        self._command_types = command_types

    def add_arguments(self, parser):
        # type: (ArgumentParser) -> None
        pass

    @contextmanager
    def parsed_command(
        self,
        args=None,  # type: Optional[Sequence[str]]
        rewrite_prog=True,  # type: bool
    ):
        # type: (...) -> Iterator[_C]
        logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)

        # By default, let argparse derive prog from sys.argv[0].
        prog = self._prog
        if os.path.basename(sys.argv[0]) == "__main__.py" and rewrite_prog:
            prog = "{python} -m {module}".format(
                python=sys.executable, module=".".join(type(self).__module__.split(".")[:-1])
            )

        parser = ArgumentParser(
            prog=prog,
            formatter_class=ArgumentDefaultsHelpFormatter,
            description=self._description,
        )
        parser.add_argument("-V", "--version", action="version", version=__version__)
        parser.set_defaults(command_type=functools.partial(Command.show_help, parser))
        self.add_arguments(parser)
        if self._command_types:
            subparsers = parser.add_subparsers(description=self._subparsers_description)
            for command_type in self._command_types:
                name = command_type.name()
                description = command_type.description()
                help_text = description.splitlines()[0] if description else None
                command_parser = subparsers.add_parser(
                    name,
                    formatter_class=ArgumentDefaultsHelpFormatter,
                    help=help_text,
                    description=description,
                )
                command_type.add_arguments(command_parser)
                command_parser.set_defaults(command_type=command_type)

        args = args or sys.argv[1:]
        passthrough_args = None  # type: Optional[Tuple[str, ...]]
        try:
            passthrough_divide = args.index("--")
        except ValueError:
            pass
        else:
            passthrough_args = tuple(args[passthrough_divide + 1 :])
            args = args[:passthrough_divide]

        options, unknown_args = parser.parse_known_args(args=args)
        with global_environment(options):
            command_type = cast("Type[_C]", options.command_type)
            if unknown_args and not command_type.supports_unknown_args():
                parser.error(
                    "Unrecognized arguments: {args}".format(
                        args=" ".join(shlex_quote(arg) for arg in unknown_args)
                    )
                )
            if passthrough_args:
                unknown_args.extend(passthrough_args)
            if unknown_args:
                passthrough_args = tuple(unknown_args)
            yield command_type(options, passthrough_args)
