"""
PIP handling for tiamat packaged python projects.
"""
import logging
import os
import pathlib
import runpy
import sys
import traceback
from codeop import CommandCompiler
from types import CodeType
from typing import cast
from typing import List
from typing import Optional
from typing import Sequence
from typing import TextIO
from typing import Union

import pip._vendor.distlib
from pip._internal.cli.main import main as pip_main
from pip._internal.commands.install import InstallCommand
from pip._internal.commands.uninstall import UninstallCommand

from tiamatpip import configure
from tiamatpip import patches
from tiamatpip.store import DistributionNotFound
from tiamatpip.store import Store
from tiamatpip.utils import debug_print
from tiamatpip.utils import patched_environ
from tiamatpip.utils import patched_sys_argv
from tiamatpip.utils import prepend_sys_path


if getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS"):
    # We're running from a compiled pyinstaller binary and distlib knows
    # nothing about this. It also does not know how to handle pyimod03_importers, so,
    # we need to patch the __loader__ attribute to force distlib to it's fallback code.
    # This is particularly important on Windows where failing to do so makes tiamat-pip
    # unusable.
    pip._vendor.distlib.__loader__ = None


# If there are logging handlers already configured, then the basicConfig
# call below will be a no-op
logging.basicConfig(
    stream=sys.stderr,
    format="%(message)s",
    level=logging.DEBUG if "TIAMAT_PIP_DEBUG" in os.environ else logging.INFO,
)
log = logging.getLogger(__name__)


class Unbuffered:
    """
    Simple wrapper class to make the passed stream unbuffered.
    """

    def __init__(self, stream):
        self.stream = stream

    def write(self, data):
        """
        Write data.
        """
        self.stream.write(data)
        self.stream.flush()

    def writelines(self, datas):
        """
        Write lines.
        """
        self.stream.writelines(datas)
        self.stream.flush()

    def __getattr__(self, attr):
        """
        Return an attribute of the wrapped stream.
        """
        return getattr(self.stream, attr)


def should_redirect_argv(argv: List[str]) -> bool:
    """
    Check if ``argv`` should be redirected.
    """
    with debug_print("should_redirect_argv", argv):
        if "TIAMAT_PIP_INSTALL" in os.environ:
            # A pip command is already in progress. This is usually
            # hit when pip is building the dependencies of a package
            log.debug(
                "argv should be redirected to pip because TIAMAT_PIP_INSTALL "
                "was found on the environ"
            )
            return True
        pip_command_name = configure.get_pip_command_name()
        if len(argv) < 2:
            log.debug(
                "argv should not be redirected to pip because argv is smaller than expected: %s",
                argv,
            )
            return False
        elif argv[1] == pip_command_name:
            # We should intercept pip comands
            log.debug(
                "argv should be redirected to pip because argv[1]==%s", pip_command_name
            )
            return True
        elif argv[1] == "tiamat-pip-shebang":
            log.debug("Being called as a shebang, process argv in tiamat-pip")
            return True
        log.debug(
            "argv should not be redirected to pip because 'argv[1](%s) != %s' and "
            "TIAMAT_PIP_INSTALL was not found in the environment",
            argv[1],
            pip_command_name,
        )
        # Do nothing
        return False


def process_argv(argv: List[str]) -> Optional[bool]:
    """
    Process ``argv`` if it looks like it should be processed.

    Does not return if argv is processed.

    If this function returns, we don't know how to execute the argv and did nothing.
    """
    # Disable any output buffering
    sys.stdout = cast(TextIO, Unbuffered(sys.stdout))
    sys.stderr = cast(TextIO, Unbuffered(sys.stderr))
    pypath = configure.get_user_base_path()
    assert pypath
    with debug_print(
        "process_argv", argv, pypath=pypath, pypath_exists=pypath.exists()
    ):
        if not should_redirect_argv(argv):
            raise Exception(
                "Please call `should_redirect_argv` before calling process_argv"
            )

        cpath = os.environ.get("CPATH") or None
        c_include_path = os.environ.get("C_INCLUDE_PATH") or None

        pyinstaller_extract_path = sys._MEIPASS  # type: ignore[attr-defined]
        included_python_headers_path = str(
            pathlib.Path(pyinstaller_extract_path).resolve() / "include" / "python"
        )
        if cpath is None:
            cpath = included_python_headers_path
        else:
            cpath_parts = cpath.split(os.pathsep)
            if included_python_headers_path not in cpath_parts:
                cpath_parts.append(included_python_headers_path)
            cpath = os.pathsep.join(cpath_parts)
        if c_include_path is None:
            c_include_path = included_python_headers_path
        else:
            c_include_path_parts = c_include_path.split(os.pathsep)
            if included_python_headers_path not in c_include_path_parts:
                c_include_path_parts.append(included_python_headers_path)
            c_include_path = os.pathsep.join(c_include_path_parts)
        with patched_environ(C_INCLUDE_PATH=c_include_path, CPATH=cpath):
            if argv[1] == "-c":
                # Example:
                #   python -c "print 'Foo!'"
                run_code(argv[2:])
                return True

            if argv[1] == "-u" and argv[2] == "-c":
                # Example:
                #   python -u -c "print 'Foo!'"
                run_code(argv[3:])
                return True

            if argv[1] == "tiamat-pip-shebang":
                with debug_print("tiamat-pip-shebang", argv):
                    run_python_file(argv[2:])
                return True

            try:
                argv0_file = pathlib.Path(argv[0]).resolve()
                if (
                    str(argv0_file) != sys.executable
                    and argv0_file.is_file()
                    and not str(argv0_file).endswith(f"{os.sep}pip")
                ):
                    # Example:
                    #   this-is-a-script.py arg1 arg2
                    run_python_file(argv[0:])
                    return True
            except ValueError:
                # Not a valid file
                pass

            try:
                argv1_file = pathlib.Path(argv[1]).resolve()
                if argv1_file.is_file() and not str(argv1_file).endswith(
                    f"{os.sep}pip"
                ):
                    # Example:
                    #   python this-is-a-script.py arg1 arg2
                    run_python_file(argv[1:])
                    return True
            except ValueError:
                # Not a valid file
                pass

            if argv[1] == "-m" and argv[2] == "pip":
                # Example:
                #   python -m pip install foo
                argv.pop(1)
            log.debug("argv before redirecting to pip: %s", argv)
            redirect_to_pip(argv)
            return True


def process_pip_argv(argv: List[str]) -> None:
    """
    Process pip ``argv``.
    """
    pypath = configure.get_user_base_path()
    assert pypath
    with debug_print(
        "process_pip_argv", argv, pypath=pypath, pypath_exists=pypath.exists()
    ):
        if pypath is None:
            raise RuntimeError(
                "You need to run 'tiamatpip.configure.set_user_base_path(<path>)' "
                "before calling tiamatpip.cli.process_pip_argv()"
            )

        if not pypath.is_dir():
            print(
                f"The path '{pypath}' does not exist or could not be created.",
                file=sys.stderr,
                flush=True,
            )
            sys.exit(1)

        environ = {
            "TIAMAT_PIP_INSTALL": "1",
            "TIAMAT_PIP_PYPATH": str(pypath),
        }
        with patched_environ(environ=environ):
            process_argv(argv)
    raise Exception("This function should never return")


def redirect_to_pip(argv: List[str]) -> None:
    """
    Redirect ``argv`` to pip.
    """
    pypath = configure.get_user_base_path()
    assert pypath
    with debug_print(
        "redirect_to_pip", argv, pypath=pypath, pypath_exists=pypath.exists()
    ):
        targets: Sequence[str] = ("install", "list", "freeze", "uninstall", "frozen")
        try:
            cmd = argv[2]
        except IndexError:
            msg: str = "Must pass in available pip command, some of which are:"
            for cmd in targets:
                msg += f"\n - {cmd}"
            msg += (
                "\n\nIf you're missing a command which is supported by pip but not tiamat-pip "
                "please open a bug report."
            )
            print(msg, file=sys.stderr, flush=True)
            sys.exit(1)

        # Valid command found

        extra_environ = {
            "PIP_DISABLE_PIP_VERSION_CHECK": "1",
            # We don't want the setuptools distutils hack to be used
            "SETUPTOOLS_USE_DISTUTILS": "builtin",
        }

        extra_cmd_info = ""
        if cmd in ("install", "uninstall"):
            include_in_store = True
        else:
            include_in_store = False

        if cmd in ("install", "uninstall"):
            args = [cmd]
        elif cmd in ("list", "freeze", "frozen"):
            if cmd == "frozen":
                cmd = "list"
                configure.unset_user_base_path()
                extra_cmd_info = " (tiamat-pip frozen command)"
            else:
                patches.patch_pip_internal_metadata_get_environment()
            args = [cmd]
        else:
            args = [cmd]
        log.debug("Running the pip command '%s'%s", cmd, extra_cmd_info)
        args.extend(argv[3:])
        pkgs = []
        parser: Union[InstallCommand, UninstallCommand]
        if cmd in ("install", "uninstall"):
            if cmd == "install":
                parser = InstallCommand("name", "summary")
            else:
                parser = UninstallCommand("name", "summary")
            _, _args = parser.parse_args(args[:])
            pkgs.extend(_args[1:])
            log.debug(f"Packages to {cmd}: {pkgs}")

        # Call pip
        with patched_environ(environ=extra_environ):
            try:
                exitcode = call_pip(args)
                if exitcode == 0 and include_in_store is True:
                    store = Store()
                    for pkg in pkgs:
                        try:
                            if cmd == "install":
                                store.add(pkg)
                            else:
                                store.remove(pkg)
                        except DistributionNotFound as exc:
                            log.debug(f"Error adding {pkg} to store: {exc}")
                    store.write()
                sys.exit(exitcode)
            finally:
                pypath = configure.get_user_base_path()
                assert pypath
                debug_print(
                    "redirect_to_pip finally",
                    args,
                    pypath=pypath,
                    pypath_exists=pypath.exists(),
                )


def call_pip(argv: List[str]) -> int:
    """
    Call ``pip``.
    """
    with debug_print("call_pip", argv):
        exitcode: int = pip_main(argv)
        return exitcode


def _run_code(source, filename=None, exec_locals=None):
    excludes = None
    if os.environ.get("TIAMAT_PIP_INSTALL", "0") == "1":
        # If by any chance, setuptools/distutils is trying to build a python
        # extension modules(ie, C code for example), let's make sure it uses
        # the tiamat-pip pypath required scheme
        patches.patch_get_scheme_distutils()
        excludes = [str(configure.get_user_site_path())]

    if exec_locals is None:
        exec_locals = {"__name__": "__console__", "__doc__": None}
    if filename is None:
        filename = "<tiamat-pip-run-code>"
    else:
        # We want scripts which have an 'if __name__ == "__main__":'
        # section to run it.
        exec_locals["__name__"] = "__main__"
        exec_locals["__file__"] = filename

    compiler = CommandCompiler()
    try:
        code: Optional[CodeType] = compiler(source, filename, "exec")
    except (OverflowError, SyntaxError, ValueError):
        traceback.print_exc()
        sys.exit(1)

    assert code

    with prepend_sys_path(os.getcwd(), excludes=excludes):
        _handle_runpy_return(runpy._run_code, code, exec_locals)  # type: ignore[attr-defined]


def _handle_runpy_return(runpy_f, *args, **kwargs):
    try:
        with debug_print(f"_handle_runpy_return({runpy_f.__name__})", sys.argv):
            runpy_f(*args, **kwargs)
    except SystemExit as exc:
        if exc.code is not None:
            if isinstance(exc.code, str):
                print(exc.code, file=sys.stderr, flush=True)
                sys.exit(1)
            sys.exit(exc.code)
        sys.exit(1)
    except Exception:  # noqa: E722
        traceback.print_exc()
        sys.exit(1)
    sys.exit(0)


def run_code(argv: List[str]) -> None:
    """
    Run a code string, emulating `python -c "..."`.

    Does not return.
    """
    source, *_ = argv
    with debug_print("run_code", argv):
        with patched_sys_argv(argv):
            _run_code(source)


def run_python_file(argv: List[str]) -> None:
    """
    Run a python file, emulating `python file.py`.

    Does not return.
    """
    with debug_print("run_python_file", argv):
        python_file = argv[0]
        with patched_sys_argv(argv):
            _handle_runpy_return(runpy.run_path, python_file, run_name="__main__")
