import compileall
import os
import shutil
import subprocess
import sys
import sysconfig
import textwrap
import venv as _venv
from pathlib import Path
from typing import TYPE_CHECKING, Dict, Optional, Union

import virtualenv as _virtualenv

if TYPE_CHECKING:
    # Literal was introduced in Python 3.8.
    from typing import Literal

    VirtualEnvironmentType = Literal["virtualenv", "venv"]
else:
    VirtualEnvironmentType = str


class VirtualEnvironment:
    """
    An abstraction around virtual environments, currently it only uses
    virtualenv but in the future it could use pyvenv.
    """

    def __init__(
        self,
        location: Path,
        template: Optional["VirtualEnvironment"] = None,
        venv_type: Optional[VirtualEnvironmentType] = None,
    ) -> None:
        self.location = location
        assert template is None or venv_type is None
        self._venv_type: VirtualEnvironmentType
        if template is not None:
            self._venv_type = template._venv_type
        elif venv_type is not None:
            self._venv_type = venv_type
        else:
            self._venv_type = "virtualenv"
        self._user_site_packages = False
        self._template = template
        self._sitecustomize: Optional[str] = None
        self._update_paths()
        self._create()

    @property
    def _legacy_virtualenv(self) -> bool:
        if self._venv_type != "virtualenv":
            return False
        return int(_virtualenv.__version__.split(".", 1)[0]) < 20

    def __update_paths_legacy(self) -> None:
        home, lib, inc, bin = _virtualenv.path_locations(self.location)
        self.bin = Path(bin)
        self.site = Path(lib) / "site-packages"
        # Workaround for https://github.com/pypa/virtualenv/issues/306
        if hasattr(sys, "pypy_version_info"):
            version_dir = str(sys.version_info.major)
            self.lib = Path(home, "lib-python", version_dir)
        else:
            self.lib = Path(lib)

    def _update_paths(self) -> None:
        if self._legacy_virtualenv:
            self.__update_paths_legacy()
            return
        bases = {
            "installed_base": self.location,
            "installed_platbase": self.location,
            "base": self.location,
            "platbase": self.location,
        }
        paths = sysconfig.get_paths(vars=bases)
        self.bin = Path(paths["scripts"])
        self.site = Path(paths["purelib"])
        self.lib = Path(paths["stdlib"])

    def __repr__(self) -> str:
        return f"<VirtualEnvironment {self.location}>"

    def _create(self, clear: bool = False) -> None:
        if clear:
            shutil.rmtree(self.location)
        if self._template:
            # On Windows, calling `_virtualenv.path_locations(target)`
            # will have created the `target` directory...
            if (
                self._legacy_virtualenv
                and sys.platform == "win32"
                and self.location.exists()
            ):
                self.location.rmdir()
            # Clone virtual environment from template.
            shutil.copytree(self._template.location, self.location, symlinks=True)
            self._sitecustomize = self._template.sitecustomize
            self._user_site_packages = self._template.user_site_packages
        else:
            # Create a new virtual environment.
            if self._legacy_virtualenv:
                subprocess.check_call(
                    [
                        sys.executable,
                        "-m",
                        "virtualenv",
                        "--no-pip",
                        "--no-wheel",
                        "--no-setuptools",
                        os.fspath(self.location),
                    ]
                )
                self._fix_legacy_virtualenv_site_module()
            elif self._venv_type == "virtualenv":
                _virtualenv.cli_run(
                    [
                        "--no-pip",
                        "--no-wheel",
                        "--no-setuptools",
                        os.fspath(self.location),
                    ],
                )
            elif self._venv_type == "venv":
                builder = _venv.EnvBuilder()
                context = builder.ensure_directories(self.location)
                builder.create_configuration(context)
                builder.setup_python(context)
                self.site.mkdir(parents=True, exist_ok=True)
            else:
                raise RuntimeError(f"Unsupported venv type {self._venv_type!r}")
            self.sitecustomize = self._sitecustomize
            self.user_site_packages = self._user_site_packages

    def _fix_legacy_virtualenv_site_module(self) -> None:
        # Patch `site.py` so user site work as expected.
        site_py = self.lib / "site.py"
        with open(site_py) as fp:
            site_contents = fp.read()
        for pattern, replace in (
            (
                # Ensure enabling user site does not result in adding
                # the real site-packages' directory to `sys.path`.
                ("\ndef virtual_addsitepackages(known_paths):\n"),
                (
                    "\ndef virtual_addsitepackages(known_paths):\n"
                    "    return known_paths\n"
                ),
            ),
            (
                # Fix sites ordering: user site must be added before system.
                (
                    "\n    paths_in_sys = addsitepackages(paths_in_sys)"
                    "\n    paths_in_sys = addusersitepackages(paths_in_sys)\n"
                ),
                (
                    "\n    paths_in_sys = addusersitepackages(paths_in_sys)"
                    "\n    paths_in_sys = addsitepackages(paths_in_sys)\n"
                ),
            ),
        ):
            assert pattern in site_contents
            site_contents = site_contents.replace(pattern, replace)
        with open(site_py, "w") as fp:
            fp.write(site_contents)
        # Make sure bytecode is up-to-date too.
        assert compileall.compile_file(str(site_py), quiet=1, force=True)

    def _customize_site(self) -> None:
        if self._legacy_virtualenv:
            contents = ""
        else:
            # Enable user site (before system).
            contents = textwrap.dedent(
                f"""
                import os, site, sys
                if not os.environ.get('PYTHONNOUSERSITE', False):
                    site.ENABLE_USER_SITE = {self._user_site_packages}
                    # First, drop system-sites related paths.
                    original_sys_path = sys.path[:]
                    known_paths = set()
                    for path in site.getsitepackages():
                        site.addsitedir(path, known_paths=known_paths)
                    system_paths = sys.path[len(original_sys_path):]
                    for path in system_paths:
                        if path in original_sys_path:
                            original_sys_path.remove(path)
                    sys.path = original_sys_path
                    # Second, add user-site.
                    if {self._user_site_packages}:
                        site.addsitedir(site.getusersitepackages())
                    # Third, add back system-sites related paths.
                    for path in site.getsitepackages():
                        site.addsitedir(path)
                """
            ).strip()
        if self._sitecustomize is not None:
            contents += "\n" + self._sitecustomize
        sitecustomize = self.site / "sitecustomize.py"
        sitecustomize.write_text(contents)
        # Make sure bytecode is up-to-date too.
        assert compileall.compile_file(str(sitecustomize), quiet=1, force=True)

    def _rewrite_pyvenv_cfg(self, replacements: Dict[str, str]) -> None:
        pyvenv_cfg = self.location.joinpath("pyvenv.cfg")
        lines = pyvenv_cfg.read_text(encoding="utf-8").splitlines()

        def maybe_replace_line(line: str) -> str:
            key = line.split("=", 1)[0].strip()
            try:
                value = replacements[key]
            except KeyError:  # No need to replace.
                return line
            return f"{key} = {value}"

        lines = [maybe_replace_line(line) for line in lines]
        pyvenv_cfg.write_text("\n".join(lines), encoding="utf-8")

    def clear(self) -> None:
        self._create(clear=True)

    def move(self, location: Union[Path, str]) -> None:
        shutil.move(os.fspath(self.location), location)
        self.location = Path(location)
        self._update_paths()

    @property
    def sitecustomize(self) -> Optional[str]:
        return self._sitecustomize

    @sitecustomize.setter
    def sitecustomize(self, value: str) -> None:
        self._sitecustomize = value
        self._customize_site()

    @property
    def user_site_packages(self) -> bool:
        return self._user_site_packages

    @user_site_packages.setter
    def user_site_packages(self, value: bool) -> None:
        self._user_site_packages = value
        if self._legacy_virtualenv:
            marker = self.lib / "no-global-site-packages.txt"
            if self._user_site_packages:
                marker.unlink()
            else:
                marker.touch()
        else:
            self._rewrite_pyvenv_cfg(
                {"include-system-site-packages": str(bool(value)).lower()}
            )
            self._customize_site()
