"""Run performance comparisons for versions of coverage"""

import contextlib
import dataclasses
import os
import shutil
import statistics
import subprocess
import time
from pathlib import Path

from typing import Iterator, List, Optional, Tuple


class ShellSession:
    """A logged shell session.

    The duration of the last command is available as .last_duration.
    """

    def __init__(self, output_filename: str):
        self.output_filename = output_filename
        self.last_duration: float = 0
        self.foutput = None

    def __enter__(self):
        self.foutput = open(self.output_filename, "a", encoding="utf-8")
        print(f"Logging output to {os.path.abspath(self.output_filename)}")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.foutput.close()

    def print(self, *args, **kwargs):
        """Print a message to this shell's log."""
        print(*args, **kwargs, file=self.foutput)

    def run_command(self, cmd: str) -> str:
        """
        Run a command line (with a shell).

        Returns:
            str: the output of the command.

        """
        self.print(f"\n========================\n$ {cmd}")
        start = time.perf_counter()
        proc = subprocess.run(
            cmd,
            shell=True,
            check=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        output = proc.stdout.decode("utf-8")
        self.last_duration = time.perf_counter() - start
        self.print(output, end="")
        self.print(f"(was: {cmd})")
        self.print(f"(in {os.getcwd()}, duration: {self.last_duration:.3f}s)")

        if proc.returncode != 0:
            self.print(f"ERROR: command returned {proc.returncode}")
            raise Exception(
                f"Command failed ({proc.returncode}): {cmd!r}, output was:\n{output}"
            )

        return output.strip()


def rmrf(path: Path) -> None:
    """
    Remove a directory tree.  It's OK if it doesn't exist.
    """
    if path.exists():
        shutil.rmtree(path)


@contextlib.contextmanager
def change_dir(newdir: Path) -> Iterator[Path]:
    """
    Change to a new directory, and then change back.

    Will make the directory if needed.
    """
    old_dir = os.getcwd()
    newdir.mkdir(parents=True, exist_ok=True)
    os.chdir(newdir)
    try:
        yield newdir
    finally:
        os.chdir(old_dir)


@contextlib.contextmanager
def file_replace(file_name: Path, old_text: str, new_text: str) -> Iterator[None]:
    """
    Replace some text in `file_name`, and change it back.
    """
    if old_text:
        file_text = file_name.read_text()
        if old_text not in file_text:
            raise Exception("Old text {old_text!r} not found in {file_name}")
        updated_text = file_text.replace(old_text, new_text)
        file_name.write_text(updated_text)
    try:
        yield
    finally:
        if old_text:
            file_name.write_text(file_text)


class ProjectToTest:
    """Information about a project to use as a test case."""

    # Where can we clone the project from?
    git_url: Optional[str] = None

    def __init__(self):
        if self.git_url:
            self.slug = self.git_url.split("/")[-1]
            self.dir = Path(self.slug)

    def get_source(self, shell):
        """Get the source of the project."""
        if self.dir.exists():
            rmrf(self.dir)
        shell.run_command(f"git clone {self.git_url}")

    def prep_environment(self, env):
        """Prepare the environment to run the test suite.

        This is not timed.
        """
        pass

    def run_no_coverage(self, env):
        """Run the test suite with no coverage measurement."""
        pass

    def run_with_coverage(self, env, pip_args, cov_options):
        """Run the test suite with coverage measurement."""
        pass


class ToxProject(ProjectToTest):
    """A project using tox to run the test suite."""

    def prep_environment(self, env):
        env.shell.run_command(f"{env.python} -m pip install tox")
        self.run_tox(env, env.pyver.toxenv, "--notest")

    def run_tox(self, env, toxenv, toxargs=""):
        """Run a tox command. Return the duration."""
        env.shell.run_command(f"{env.python} -m tox -e {toxenv} {toxargs}")
        return env.shell.last_duration

    def run_no_coverage(self, env):
        return self.run_tox(env, env.pyver.toxenv, "--skip-pkg-install")

    def run_with_coverage(self, env, pip_args, cov_options):
        self.run_tox(env, env.pyver.toxenv, "--notest")
        env.shell.run_command(
            f".tox/{env.pyver.toxenv}/bin/python -m pip install {pip_args}"
        )
        return self.run_tox(env, env.pyver.toxenv, "--skip-pkg-install")


class ProjectPytestHtml(ToxProject):
    """pytest-dev/pytest-html"""

    git_url = "https://github.com/pytest-dev/pytest-html"

    def run_with_coverage(self, env, pip_args, cov_options):
        covenv = env.pyver.toxenv + "-cov"
        self.run_tox(env, covenv, "--notest")
        env.shell.run_command(f".tox/{covenv}/bin/python -m pip install {pip_args}")
        if cov_options:
            replace = ("# reference: https", f"[run]\n{cov_options}\n#")
        else:
            replace = ("", "")
        with file_replace(Path(".coveragerc"), *replace):
            env.shell.run_command("cat .coveragerc")
            env.shell.run_command(f".tox/{covenv}/bin/python -m coverage debug sys")
            return self.run_tox(env, covenv, "--skip-pkg-install")


class ProjectDateutil(ToxProject):
    """dateutil/dateutil"""

    git_url = "https://github.com/dateutil/dateutil"

    def prep_environment(self, env):
        super().prep_environment(env)
        env.shell.run_command(f"{env.python} updatezinfo.py")

    def run_no_coverage(self, env):
        env.shell.run_command("echo No option to run without coverage")
        return 0


class ProjectAttrs(ToxProject):
    """python-attrs/attrs"""

    git_url = "https://github.com/python-attrs/attrs"


class AdHocProject(ProjectToTest):
    """A standalone program to run locally."""

    def __init__(self, python_file, cur_dir=None, pip_args=None):
        super().__init__()
        self.python_file = Path(python_file)
        if not self.python_file.exists():
            raise ValueError(f"Couldn't find {self.python_file} to run ad-hoc.")
        self.cur_dir = Path(cur_dir or self.python_file.parent)
        if not self.cur_dir.exists():
            raise ValueError(f"Couldn't find {self.cur_dir} to run in.")
        self.pip_args = pip_args
        self.slug = self.python_file.name

    def get_source(self, shell):
        pass

    def prep_environment(self, env):
        env.shell.run_command(f"{env.python} -m pip install {self.pip_args}")

    def run_no_coverage(self, env):
        with change_dir(self.cur_dir):
            env.shell.run_command(f"{env.python} {self.python_file}")
        return env.shell.last_duration

    def run_with_coverage(self, env, pip_args, cov_options):
        env.shell.run_command(f"{env.python} -m pip install {pip_args}")
        with change_dir(self.cur_dir):
            env.shell.run_command(
                f"{env.python} -m coverage run {self.python_file}"
            )
        return env.shell.last_duration


class SlipcoverBenchmark(AdHocProject):
    """
    For running code from the Slipcover benchmarks.

    Clone https://github.com/plasma-umass/slipcover to /src/slipcover

    """
    def __init__(self, python_file):
        super().__init__(
            python_file=f"/src/slipcover/benchmarks/{python_file}",
            cur_dir="/src/slipcover",
            pip_args="six pyperf",
        )

class PyVersion:
    """A version of Python to use."""

    # The command to run this Python
    command: str
    # The tox environment to run this Python
    toxenv: str


class Python(PyVersion):
    """A version of CPython to use."""

    def __init__(self, major, minor):
        self.command = f"python{major}.{minor}"
        self.toxenv = f"py{major}{minor}"


class PyPy(PyVersion):
    """A version of PyPy to use."""

    def __init__(self, major, minor):
        self.command = f"pypy{major}.{minor}"
        self.toxenv = f"pypy{major}{minor}"


@dataclasses.dataclass
class Env:
    """An environment to run a test suite in."""

    pyver: PyVersion
    python: Path
    shell: ShellSession


def run_experiments(
    py_versions: List[PyVersion],
    cov_versions: List[Tuple[str, Optional[str], Optional[str]]],
    projects: List[ProjectToTest],
    num_runs=3,
):
    """Run test suites under different conditions."""

    results = []
    for proj in projects:
        print(f"Testing with {proj.slug}")
        with ShellSession(f"output_{proj.slug}.log") as shell:
            proj.get_source(shell)

            for pyver in py_versions:
                print(f"Making venv for {proj.slug} {pyver.command}")
                venv_dir = f"venv_{proj.slug}_{pyver.command}"
                shell.run_command(f"{pyver.command} -m venv {venv_dir}")
                python = Path.cwd() / f"{venv_dir}/bin/python"
                shell.run_command(f"{python} -V")
                env = Env(pyver, python, shell)

                with change_dir(Path(proj.slug)):
                    print(f"Prepping for {proj.slug} {pyver.command}")
                    proj.prep_environment(env)
                    for cov_slug, cov_pip, cov_options in cov_versions:
                        durations = []
                        for run_num in range(num_runs):
                            print(
                                f"Running tests, cov={cov_slug}, {run_num+1} of {num_runs}"
                            )
                            if cov_pip is None:
                                dur = proj.run_no_coverage(env)
                            else:
                                dur = proj.run_with_coverage(env, cov_pip, cov_options)
                            print(f"Tests took {dur:.3f}s")
                            durations.append(dur)
                        med = statistics.median(durations)
                        result = (
                            f"Median for {proj.slug}, {pyver.command}, "
                            + f"cov={cov_slug}: {med:.3f}s"
                        )
                        print(f"## {result}")
                        results.append(result)

    print("# Results")
    for result in results:
        print(result)


PERF_DIR = Path("/tmp/covperf")


print(f"Removing and re-making {PERF_DIR}")
rmrf(PERF_DIR)

with change_dir(PERF_DIR):

    if 1:
        run_experiments(
            py_versions=[
                Python(3, 10),
                Python(3, 11),
            ],
            cov_versions=[
                ("none", None, None),
                ("6.4.1", "coverage==6.4.1", ""),
                # (
                #     "tip timid",
                #     "git+https://github.com/nedbat/coveragepy.git@master",
                #     "timid=True",
                # ),
            ],
            projects=[
                # ProjectPytestHtml(),
                #ProjectAttrs(),
                AdHocProject("/src/bugs/bug1339/bug1339.py"),
                SlipcoverBenchmark("bm_sudoku.py"),
                SlipcoverBenchmark("bm_spectral_norm.py"),
            ],
            num_runs=3,
        )

    if 0:
        run_experiments(
            py_versions=[
                PyPy(3, 9),
            ],
            cov_versions=[
                ("none", None, None),
                ("6.4", "coverage==6.4", ""),
                (
                    "PR 1381",
                    "git+https://github.com/cfbolz/coveragepy.git@f_trace_lines",
                    "",
                ),
            ],
            projects=[
                ProjectPytestHtml(),
            ],
            num_runs=3,
        )

    if 0:
        run_experiments(
            py_versions=[
                PyPy(3, 9),
            ],
            cov_versions=[
                ("none", None, None),
                ("6.4", "coverage", ""),
                ("tip", "git+https://github.com/nedbat/coveragepy.git@master", ""),
            ],
            projects=[
                AdHocProject("/src/bugs/bug1339/bug1339.py"),
            ],
            num_runs=7,
        )