"""A central, mostly package-agnostic collection of Nox utilities.

This module defines the SessionBuilder class, which takes in some configuration
options and exposes methods that generate Nox sessions. See its documentation
for details.

Most sessions generated by SessionBuilder are designed to work either directly
in a development environment (i.e. with nox's --no-venv option) or in a
nox-managed virtualenv (as they would run in the CI). Sessions that only work in
one or the other will indicate this in their docstrings.
"""


import subprocess
from functools import wraps
from typing import Dict, Any, List, Optional
from pathlib import Path
import os
import re
import datetime
import nox
import platform

# The distinction between these two session types is that poetry_session
# automatically limits installations to the version numbers in the Poetry lock
# file, while nox_session does not. Otherwise, their interfaces should be
# identical.
from nox import session as nox_session
from nox_poetry import session as poetry_session

from .dependencies import install, show_installed
from .environment import in_ci, with_clean_workdir

# TODO(#2140): Once support for is added to nox-poetry (see
#   https://github.com/cjolowicz/nox-poetry/issues/663), some of the
#   installation lists here can be rewritten in terms of dependency groups,
#   making the pyproject file more of a single source for information about
#   dependencies.

class SessionBuilder:
    """Class for creating common Nox sessions based on project-specific configuration.

    Options that apply to all packages and sessions are keyword arguments to the
    constructor, while more use-case-specific options are passed as key-value
    pairs via the `options` parameter. The appropriate names and formats of
    options are given by the individual sessions that use them.

    To add a session to a particular noxfile, just create an instance of this
    class and then call the method(s) corresponding to the session(s) you want.

    This class leverages the fact that *any* function that is defined with the
    `@session` decorator is picked up as a Nox session, no matter where it is
    defined. By wrapping these definitions inside of class methods, we can
    access the class attributes to do project-specific things while still
    maintaining a common collection of sessions that can be reused across
    projects, and can avoid needing boilerplate in the noxfile itself.
    """

    def __init__(
        self, package_name: str, package_source: Path, options: Dict[str, Any]
    ):
        self.package_name = package_name
        self.package_source = package_source
        self.options = options
        self.cwd = Path(".").resolve()
        self.package_version = (
            subprocess.run(["poetry", "version", "-s"], capture_output=True)
            .stdout.decode("utf-8")
            .strip()
        )

    def _get_physical_cpu_count(self):
        """Get the number of physical CPUs on the current machine."""
        try:
            if platform.system() == "Darwin":
                command = "sysctl -n hw.physicalcpu"
            elif platform.system() == "Linux":
                command = "nproc --all"
            cores = int(subprocess.check_output(command, shell=True).strip())
        except Exception as e:
            print(f"Error getting physical CPU count: {e}, defaulting to 1")
            cores = 1
        if cores == 1:
            return cores
        return cores // 2 if in_ci() else 1

    def _build(self, session):
        """Build sdists and wheels for the package.

        If the build process for a package requires custom logic, pass the
        "build" option to SessionBuilder, with the value being function taking
        one argument (a nox Session) that builds the package. If this option is
        not passed, the package is just built with `poetry build`.
        """
        if "build" in self.options:
            session.log("Using custom build function")
            build = self.options["install_overrides"]
            assert callable(
                build
            ), "build option should be a callable taking a nox session object"
            build(session)
        else:
            session.run("poetry", "build", external=True)

    def build(self):
        @poetry_session()
        def build(session):
            """Build source distributions and wheels for this package."""
            self._build(session)

    def install_package(self, f):
        """Install the main package a dev wheel into the test virtual environment.

        Installs the package from this repository and all its dependencies, if the
        current environment supports installing packages. If wheels for
        the current dev version (from `poetry version`) are not already present in
        `dist/`, this will build them.

        Similar to the @install() decorator, this decorator automatically skips
        installation in non-sandboxed environments.
        """

        @wraps(f)
        def inner(session, *args, **kwargs):
            if session.virtualenv.is_sandboxed:
                temp_dir = session.create_tmp()
                out = session.run(
                    "pip",
                    "download",
                    f"{self.package_name}=={self.package_version}",
                    "--find-links",
                    f"{self.cwd}/dist/",
                    "--only-binary",
                    self.package_name,
                    "-d",
                    temp_dir,
                    "--no-deps",
                    silent=True,
                    success_codes=[0, 1],
                )
                if "No matching distribution" in out:
                    self._build(session)

                session.install(
                    f"{self.package_name}=={self.package_version}",
                    "--find-links",
                    f"{self.cwd}/dist/",
                    "--only-binary",
                    self.package_name,
                )
                if "install_overrides" in self.options:
                    session.log("Running install_overrides")
                    install_overrides = self.options["install_overrides"]
                    assert callable(install_overrides), (
                        "install_overrides option should be a callable "
                        "taking a nox session object"
                    )
                    install_overrides(session)
            else:
                session.log("Skipping package installation, non-sandboxed environment")
            return f(session, *args, **kwargs)

        return inner

    def _get_code_dirs(self):
        assert (
            "code_dirs" in self.options
        ), "code_dirs option must be set to list of code directories, including tests"
        code_dirs = self.options["code_dirs"]
        assert isinstance(
            code_dirs, list
        ), "code_dirs option must be set to list of code directories, including tests"
        assert len(code_dirs) > 0, "code_dirs option must not be an empty list"
        assert all(
            isinstance(d, Path) for d in code_dirs
        ), "code_dirs option elements must all be Path objects"
        return code_dirs

    def black(self):
        @poetry_session(tags=["lint"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("black")
        @show_installed
        def black(session):
            """Run black. If the --check argument is given, only check, don't make changes."""
            check_flags = ["--check", "--diff"] if "--check" in session.posargs else []
            session.run(
                "black",
                *check_flags,
                *map(str, self._get_code_dirs()),
            )

    def isort(self):
        @poetry_session(tags=["lint"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("isort[pyproject]", "pytest")
        @show_installed
        def isort(session):
            check_flags = (
                ["--check-only", "--diff"] if "--check" in session.posargs else []
            )
            session.run("isort", *check_flags, *map(str, self._get_code_dirs()))

    def mypy(self):
        @poetry_session(tags=["lint"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("mypy")
        @show_installed
        def mypy(session):
            session.run("mypy", *map(str, self._get_code_dirs()))

    def pylint(self):
        @poetry_session(tags=["lint"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("pylint", "pytest")
        @show_installed
        def pylint(session):
            session.run("pylint", "--score=no", *map(str, self._get_code_dirs()))

    def pydocstyle(self):
        @poetry_session(tags=["lint"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("pydocstyle[toml]")
        @show_installed
        def pydocstyle(session):
            session.run("pydocstyle", *map(str, self._get_code_dirs()))

    def audit(self):
        """Generate the audit parameterized Nox session.

        To be useful, this function needs the `audit_versions` option to be
        set. It should contain a list of Python versions to be audited.
        """
        python_versions = self.options.get("audit_versions")
        if not python_versions:
            raise KeyError(
                "audit_versions option must be set to a list of Python versions"
            )

        ignore_vulns = self.options.get("audit_suppressions", [])
        # Should produce --ignore-vuln [vuln_id] for each vulnerability ID
        ignore_options = [
            i
            for vuln in ignore_vulns
            for i in ("--ignore-vuln", vuln)
        ]

        @nox_session
        @self.install_package
        @install("pip-audit")
        @show_installed
        @nox.parametrize("python", python_versions)
        def audit(session):
            session.run(
                "pip-audit", "-v", "--progress-spinner", "off",
                *ignore_options
            )

    @install("pytest", "pytest-xdist", "pytest-cov")
    @show_installed
    @with_clean_workdir
    def _test(
        self,
        session,
        test_paths: Optional[List[Path]] = None,
        min_coverage: Optional[int] = None,
        extra_args: Optional[List[str]] = None,
    ):
        test_paths = test_paths or self._get_code_dirs()
        if min_coverage is None:
            min_coverage = self.options.get("minimum_coverage", 0)
        extra_args = extra_args or []
        # If the user passes args, pass them on to pytest. The main reason this is
        # useful is for specifying a particular subset of tests to run, so clear
        # test_paths to allow that use case.
        if session.posargs:
            test_paths = []
            extra_args.extend(session.posargs)
        test_options = [
            "-r fEs",
            "--verbose",
            "--disable-warnings",
            f"--junitxml={self.cwd}/junit.xml",
            # Show runtimes of the 10 slowest tests, for later comparison if needed.
            "--durations=10",
            "-n",
            f"{self._get_physical_cpu_count()}",
            # Collect coverage data, enforce minimum, output reports
            f"--cov={self.options['coverage_module']}",
            f"--cov-fail-under={min_coverage}",
            "--cov-report=term",
            f"--cov-report=html:{self.cwd}/coverage/",
            f"--cov-report=xml:{self.cwd}/coverage.xml",
            # Any extra args
            *extra_args,
            # The files to be tested
            *map(str, test_paths),
        ]
        session.run("pytest", *test_options)

    @show_installed
    @with_clean_workdir
    def _smoketest(self, session):
        smoketest_script = self.options.get("smoketest_script")
        if not smoketest_script:
            session.error("smoketest_script option must be set when running smoketest")
        session.run("python", "-c", smoketest_script)

    def test(self):
        @poetry_session(tags=["test"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def test(session):
            """Run all tests."""
            self._test(session)

    def test_doctest(self):
        @poetry_session(tags=["test"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def test_doctest(session):
            """Run documentation tests."""
            self._test(
                session,
                test_paths=[self.package_source],
                min_coverage=0,
                extra_args=["--doctest-modules"],
            )

    def test_demos(self):
        @poetry_session(tags=["test"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("notebook", "nbconvert", "matplotlib", "seaborn")
        @show_installed
        def test_demos(session):
            """Run all examples."""
            demos_path = self.cwd / "demos"
            if not demos_path.exists():
                session.error("No demos directory found, nothing to run")
            demos_py = []
            demos_ipynb = []
            unknown = []
            ignored = []
            for f in demos_path.iterdir():
                if f.is_file() and f.suffix == ".py":
                    demos_py.append(f)
                elif f.is_file() and f.suffix == ".ipynb":
                    if ".nbconvert" not in f.suffixes:
                        demos_ipynb.append(f)
                    else:
                        ignored.append(f)
                else:
                    unknown.append(f)
            for py in demos_py:
                session.run("python", str(py))
            for nb in demos_ipynb:
                session.run(
                    "jupyter", "nbconvert", "--to=notebook", "--execute", str(nb)
                )
            if ignored:
                session.log(f"Ignored: {', '.join(str(f) for f in ignored)}")
            if unknown:
                session.warn(
                    f"Found unknown files in demos: {', '.join(str(f) for f in unknown)}"
                )

    def test_smoketest(self):
        @poetry_session(tags=["test"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def test_smoketest(session):
            """Run a smoketest."""
            self._smoketest(session)

    # test_fast and test_slow don't get the "test" tag, as they run subsets of
    # the test suite, so there's no reason to run them again when using `nox -t
    # test`.

    def test_fast(self):
        @poetry_session(python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def test_fast(session):
            """Run all fast tests."""
            self._test(session, extra_args=["-m", "not slow"])

    def test_slow(self):
        @poetry_session(python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def test_slow(session):
            """Run all slow tests."""
            self._test(session, min_coverage=0, extra_args=["-m", "slow"])

    def test_dependency_matrix(self):
        """Generate the test_dependency_matrix parameterized Nox session.

        To be useful, this function needs the `dependency_matrix` option to be
        set. It should contain a dictionary whose keys are the names for the
        different configurations (which are used as the keys for the sessions),
        and whose values are dictionaries as follows:
          - There must be a "python" key, whose value is the Python minor
            version to be used for that configuration, e.g. "3.10".
          - All other keys are interpreted as package names, and their values as
            PEP440 version specifiers
            (https://peps.python.org/pep-0440/#version-specifiers) that will be
            used to install a specified version of the named package. For
            example, having `{"tmlt.core": ">=0.11.2,<0.12.0"}` will, in effect,
            run `pip install tmlt.core>=0.11.2,<0.12.0` after installing the
            main package.
        """
        dependency_matrix = self.options.get("dependency_matrix")
        if not dependency_matrix:
            return

        ids = []
        pythons = []
        packages = []
        for config_id, config in dependency_matrix.items():
            ids.append(config_id)
            try:
                pythons.append(config.pop("python"))
            except KeyError:
                raise RuntimeError(
                    "Dependency matrix configurations must specify a Python minor version"
                )
            packages.append(config)

        @nox_session
        @install("pytest", "pytest-xdist", "pytest-cov")
        @with_clean_workdir
        @nox.parametrize("python,packages", zip(pythons, packages), ids=ids)
        def test_dependency_matrix(session, packages):
            """Run tests using various dependencies."""
            session.install(
                f"{self.package_name}=={self.package_version}",
                "--find-links",
                f"{self.cwd}/dist/",
                "--only-binary",
                self.package_name,
            )
            session.install(*[pkg + version for pkg, version in packages.items()])
            session.run("pip", "freeze")

            test_selector = session.posargs or [
                "-m",
                "not slow",
                *map(str, self._get_code_dirs()),
            ]
            test_options = [
                "-rfs",
                "--disable-warnings",
                f"--junitxml={self.cwd}/junit.xml",
                # Show runtimes of the 10 slowest tests, for later comparison if needed.
                "--durations=10",
                "-n",
                f"{self._get_physical_cpu_count()}",
                # Collect coverage data, enforce minimum, output reports
                f"--cov={self.options['coverage_module']}",
                "--cov-report=term",
                f"--cov-report=html:{self.cwd}/coverage/",
                f"--cov-report=xml:{self.cwd}/coverage.xml",
                # The files to be tested
                *test_selector,
            ]
            session.run("pytest", *test_options)

    @install(
        "pandoc",
        "pydata-sphinx-theme",
        "scanpydoc",
        "sphinx",
        "sphinx-autoapi",
        "sphinx-autodoc-typehints",
        "sphinx-copybutton",
        "sphinx-panels",
        "sphinxcontrib-bibtex",
        "sphinxcontrib-images",
    )
    def _run_sphinx(self, session, builder: str):
        sphinx_options = ["-n", "-W", "--keep-going"]
        session.run("sphinx-build", "doc/", "public/", f"-b={builder}", *sphinx_options)

    def docs_linkcheck(self):
        @poetry_session(tags=["docs"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def docs_linkcheck(session):
            """Run linkcheck on docs."""
            self._run_sphinx(session, "linkcheck")

    def docs_doctest(self):
        @poetry_session(tags=["docs"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        @install("matplotlib", "seaborn")
        def docs_doctest(session):
            """Run doctest on code examples in documentation."""
            self._run_sphinx(session, "doctest")

    def docs(self):
        @poetry_session(tags=["docs"], python="3.7")  # type: ignore[call-overload]
        @self.install_package
        def docs(session):
            """Generate HTML documentation."""
            self._run_sphinx(session, "html")

    def release_test(self):
        @nox_session()
        @self.install_package
        def release_test(session):
            """Test a wheel as it would be installed on a user's machine.

            This session is used to verify that built wheels install correctly as a user
            would install them, without Poetry. It installs a wheel given as a
            positional argument, then runs the fast tests on it.

            Note: This session doesn't do anything useful when run with the `--no-venv`
                  option, as it requires a clean environment to install things in.
            """
            self._test(session, extra_args=["-m", "not slow"], show_installed=True)

    def release_smoketest(self):
        @nox_session()
        @self.install_package
        def release_smoketest(session):
            """Smoke-test a wheel as it would be installed on a user's machine.

            This session installs a built wheel as the user would install it, without
            Poetry, then runs a short test to ensure that the library plausibly works.

            Note: This session doesn't do anything useful when run with the `--no-venv`
                  option, as it requires a clean environment to install things in.
            """
            self._smoketest(session)

    def prepare_release(self):
        @nox_session(python=None)
        def prepare_release(session):
            """Update files in preparation for a release.

            The version number for the new release should be in the VERSION environment
            variable.
            """
            version = os.environ.get("VERSION", "")
            if not version:
                session.error("VERSION not set, unable to prepare release")

            # Check version number against our allowed version format. This matches a
            # subset of semantic versions that closely matches PEP440 versions. Some
            # examples include: 0.1.2, 1.2.3-alpha.2, 1.3.0-rc.1
            version_regex = (
                r"^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)"
                r"(-(alpha|beta|rc)\.(0|[1-9]\d*))?$"
            )
            if not re.match(version_regex, version):
                session.error(f"VERSION {version} is not a valid version number.")
            session.debug(f"Preparing release {version}")

            # Replace "Unreleased" section header in changelog for non-prerelease
            # releases. Between the base version and prerelease number is the only place
            # a hyphen can appear in the version number, so just checking for that
            # indicates whether a version is a prerelease.
            is_pre_release = "-" in version
            if not is_pre_release:
                session.log("Updating CHANGELOG.rst unreleased version...")
                with Path("CHANGELOG.rst").open("r", encoding="utf-8") as fp:
                    changelog_content = fp.readlines()
                for i, content in enumerate(changelog_content):
                    if re.match("^Unreleased$", content):
                        # BEFORE
                        # Unreleased
                        # ----------

                        # AFTER
                        # .. _v1.2.3:
                        #
                        # 1.2.3 - 2020-01-01
                        # ------------------
                        version_header = f"{version} - {datetime.date.today()}"
                        # The anchor is necessary because sphinx doesn't support anchors
                        # starting with numerics, so we have to generate our own
                        # reference link.
                        anchor = f".. _v{version}:\n\n"
                        subsection = f"{version_header}\n{'-' * len(version_header)}\n"
                        changelog_content[i] = anchor
                        changelog_content[i + 1] = subsection
                        break
                else:
                    session.error(
                        "Renaming unreleased section in changelog failed, "
                        "unable to find matching line"
                    )
                with Path("CHANGELOG.rst").open("w", encoding="utf-8") as fp:
                    fp.writelines(changelog_content)
            else:
                session.log("Prerelease, skipping CHANGELOG.rst update...")

    def post_release(self):
        @nox_session(python=None)
        def post_release(session):
            """Update files after a release."""
            unreleased_header = ["Unreleased\n", "----------\n", "\n"]

            anchor_regex = (
                r"^\.\. _v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)"
                r"(-(alpha|beta|rc)\.(0|[1-9]\d*))?:$"
            )
            # Find the latest release
            with Path("CHANGELOG.rst").open("r", encoding="utf-8") as fp:
                changelog_content = fp.readlines()
                for i, content in enumerate(changelog_content):
                    # If prerelease, the Unreleased header should still remain and we
                    # will skip the update.
                    if changelog_content[i : i + 3] == unreleased_header:
                        session.log(
                            "Unreleased section found, skipping CHANGELOG.rst update..."
                        )
                        return
                    if re.match(anchor_regex, content):
                        # BEFORE
                        # .. _v1.2.3:
                        #
                        # 1.2.3 - 2020-01-01
                        # ------------------

                        # AFTER
                        # Unreleased
                        # ----------
                        #
                        # .. _v1.2.3:
                        #
                        # 1.2.3 - 2020-01-01
                        # ------------------
                        for new_line in reversed(unreleased_header):
                            changelog_content.insert(i, new_line)
                        break
                else:
                    session.error("Unable to find latest release in CHANGELOG.rst")
                with Path("CHANGELOG.rst").open("w", encoding="utf-8") as fp:
                    fp.writelines(changelog_content)
