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

from __future__ import annotations

import configparser
from dataclasses import dataclass
from enum import Enum
from io import StringIO
from pathlib import PurePath
from typing import Any, MutableMapping, cast

import toml

from pants.backend.python.goals import lockfile
from pants.backend.python.goals.lockfile import GeneratePythonLockfile
from pants.backend.python.subsystems.python_tool_base import PythonToolBase
from pants.backend.python.target_types import ConsoleScript
from pants.backend.python.util_rules.pex import PexRequest, VenvPex, VenvPexProcess
from pants.backend.python.util_rules.python_sources import (
    PythonSourceFiles,
    PythonSourceFilesRequest,
)
from pants.core.goals.generate_lockfiles import GenerateToolLockfileSentinel
from pants.core.goals.test import (
    ConsoleCoverageReport,
    CoverageData,
    CoverageDataCollection,
    CoverageReport,
    CoverageReports,
    FilesystemCoverageReport,
)
from pants.core.util_rules.config_files import ConfigFiles, ConfigFilesRequest
from pants.engine.addresses import Address
from pants.engine.fs import (
    EMPTY_DIGEST,
    AddPrefix,
    CreateDigest,
    Digest,
    DigestContents,
    FileContent,
    MergeDigests,
    PathGlobs,
    Snapshot,
)
from pants.engine.process import FallibleProcessResult, ProcessExecutionFailure, ProcessResult
from pants.engine.rules import Get, MultiGet, collect_rules, rule
from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest
from pants.engine.unions import UnionRule
from pants.option.global_options import ProcessCleanupOption
from pants.option.option_types import (
    BoolOption,
    EnumListOption,
    FileOption,
    FloatOption,
    StrListOption,
    StrOption,
)
from pants.source.source_root import AllSourceRoots
from pants.util.docutil import git_url
from pants.util.logging import LogLevel

"""
An overview:

Step 1: Run each test with the appropriate `--cov` arguments.
In `python_test_runner.py`, we pass options so that the pytest-cov plugin runs and records which
lines were encountered in the test. For each test, it will save a `.coverage` file (SQLite DB
format).

Step 2: Merge the results with `coverage combine`.
We now have a bunch of individual `PytestCoverageData` values, each with their own `.coverage` file.
We run `coverage combine` to convert this into a single `.coverage` file.

Step 3: Generate the report with `coverage {html,xml,console}`.
All the files in the single merged `.coverage` file are still stripped, and we want to generate a
report with the source roots restored. Coverage requires that the files it's reporting on be present
when it generates the report, so we populate all the source files.

Step 4: `test.py` outputs the final report.
"""


class CoverageReportType(Enum):
    CONSOLE = ("console", "report")
    XML = ("xml", None)
    HTML = ("html", None)
    RAW = ("raw", None)
    JSON = ("json", None)

    _report_name: str

    def __new__(cls, value: str, report_name: str | None = None) -> CoverageReportType:
        member: CoverageReportType = object.__new__(cls)
        member._value_ = value
        member._report_name = report_name if report_name is not None else value
        return member

    @property
    def report_name(self) -> str:
        return self._report_name

    @property
    def value(self) -> str:
        return cast(str, super().value)


class CoverageSubsystem(PythonToolBase):
    options_scope = "coverage-py"
    help = "Configuration for Python test coverage measurement."

    default_version = "coverage[toml]>=5.5,<5.6"
    default_main = ConsoleScript("coverage")

    register_interpreter_constraints = True
    default_interpreter_constraints = ["CPython>=3.6,<4"]

    register_lockfile = True
    default_lockfile_resource = ("pants.backend.python.subsystems", "coverage_py_lockfile.txt")
    default_lockfile_path = "src/python/pants/backend/python/subsystems/coverage_py_lockfile.txt"
    default_lockfile_url = git_url(default_lockfile_path)

    filter = StrListOption(
        "--filter",
        help=(
            "A list of Python modules or filesystem paths to use in the coverage report, e.g. "
            "`['helloworld_test', 'helloworld/util/dirutil'].\n\nBoth modules and directory "
            "paths are recursive: any submodules or child paths, respectively, will be "
            "included.\n\nIf you leave this off, the coverage report will include every file "
            "in the transitive closure of the address/file arguments; for example, `test ::` "
            "will include every Python file in your project, whereas "
            "`test project/app_test.py` will include `app_test.py` and any of its transitive "
            "dependencies."
        ),
    )
    reports = EnumListOption(
        "--report",
        default=[CoverageReportType.CONSOLE],
        help="Which coverage report type(s) to emit.",
    )
    _output_dir = StrOption(
        "--output-dir",
        default=str(PurePath("dist", "coverage", "python")),
        advanced=True,
        help="Path to write the Pytest Coverage report to. Must be relative to build root.",
    )
    config = FileOption(
        "--config",
        default=None,
        advanced=True,
        help=lambda cls: (
            "Path to an INI or TOML config file understood by coverage.py "
            "(https://coverage.readthedocs.io/en/stable/config.html).\n\n"
            f"Setting this option will disable `[{cls.options_scope}].config_discovery`. Use "
            f"this option if the config is located in a non-standard location."
        ),
    )
    config_discovery = BoolOption(
        "--config-discovery",
        default=True,
        advanced=True,
        help=lambda cls: (
            "If true, Pants will include any relevant config files during runs "
            "(`.coveragerc`, `setup.cfg`, `tox.ini`, and `pyproject.toml`)."
            f"\n\nUse `[{cls.options_scope}].config` instead if your config is in a "
            f"non-standard location."
        ),
    )
    global_report = BoolOption(
        "--global-report",
        default=False,
        help=(
            "If true, Pants will generate a global coverage report.\n\nThe global report will "
            "include all Python source files in the workspace and not just those depended on "
            "by the tests that were run."
        ),
    )
    fail_under = FloatOption(
        "--fail-under",
        default=None,
        help=(
            "Fail if the total combined coverage percentage for all tests is less than this "
            "number.\n\nUse this instead of setting fail_under in a coverage.py config file, "
            "as the config will apply to each test separately, while you typically want this "
            "to apply to the combined coverage for all tests run."
            "\n\nNote that you must generate at least one (non-raw) coverage report for this "
            "check to trigger.\n\nNote also that if you specify a non-integral value, you must "
            "also set [report] precision properly in the coverage.py config file to make use "
            "of the decimal places. See https://coverage.readthedocs.io/en/latest/config.html ."
        ),
    )

    @property
    def output_dir(self) -> PurePath:
        return PurePath(self._output_dir)

    @property
    def config_request(self) -> ConfigFilesRequest:
        # Refer to https://coverage.readthedocs.io/en/stable/config.html.
        return ConfigFilesRequest(
            specified=self.config,
            specified_option_name=f"[{self.options_scope}].config",
            discovery=self.config_discovery,
            check_existence=[".coveragerc"],
            check_content={
                "setup.cfg": b"[coverage:",
                "tox.ini": b"[coverage:]",
                "pyproject.toml": b"[tool.coverage",
            },
        )


class CoveragePyLockfileSentinel(GenerateToolLockfileSentinel):
    resolve_name = CoverageSubsystem.options_scope


@rule
def setup_coverage_lockfile(
    _: CoveragePyLockfileSentinel, coverage: CoverageSubsystem
) -> GeneratePythonLockfile:
    return GeneratePythonLockfile.from_tool(coverage)


@dataclass(frozen=True)
class PytestCoverageData(CoverageData):
    address: Address
    digest: Digest


class PytestCoverageDataCollection(CoverageDataCollection):
    element_type = PytestCoverageData


@dataclass(frozen=True)
class CoverageConfig:
    digest: Digest
    path: str


def _validate_and_update_config(
    coverage_config: configparser.ConfigParser, config_path: str | None
) -> None:
    if not coverage_config.has_section("run"):
        coverage_config.add_section("run")
    run_section = coverage_config["run"]
    relative_files_str = run_section.get("relative_files", "True")
    if relative_files_str.lower() != "true":
        raise ValueError(
            "relative_files under the 'run' section must be set to True in the config "
            f"file {config_path}"
        )
    coverage_config.set("run", "relative_files", "True")
    omit_elements = list(run_section.get("omit", "").split("\n")) or ["\n"]
    if "pytest.pex/*" not in omit_elements:
        omit_elements.append("pytest.pex/*")
    run_section["omit"] = "\n".join(omit_elements)


class InvalidCoverageConfigError(Exception):
    pass


def _parse_toml_config(fc: FileContent) -> MutableMapping[str, Any]:
    try:
        return toml.loads(fc.content.decode())
    except toml.TomlDecodeError as exc:
        raise InvalidCoverageConfigError(
            f"Failed to parse the coverage.py config `{fc.path}` as TOML. Please either fix "
            f"the config or update `[coverage-py].config` and/or "
            f"`[coverage-py].config_discovery`.\n\nParse error: {repr(exc)}"
        )


def _parse_ini_config(fc: FileContent) -> configparser.ConfigParser:
    cp = configparser.ConfigParser()
    try:
        cp.read_string(fc.content.decode())
        return cp
    except configparser.Error as exc:
        raise InvalidCoverageConfigError(
            f"Failed to parse the coverage.py config `{fc.path}` as INI. Please either fix "
            f"the config or update `[coverage-py].config` and/or `[coverage-py].config_discovery`."
            f"\n\nParse error: {repr(exc)}"
        )


def _update_config(fc: FileContent) -> FileContent:
    if PurePath(fc.path).suffix == ".toml":
        all_config = _parse_toml_config(fc)
        tool = all_config.setdefault("tool", {})
        coverage = tool.setdefault("coverage", {})
        run = coverage.setdefault("run", {})
        run["relative_files"] = True
        if "pytest.pex/*" not in run.get("omit", []):
            run["omit"] = [*run.get("omit", []), "pytest.pex/*"]
        return FileContent(fc.path, toml.dumps(all_config).encode())

    cp = _parse_ini_config(fc)
    run_section = "coverage:run" if fc.path in ("tox.ini", "setup.cfg") else "run"
    if not cp.has_section(run_section):
        cp.add_section(run_section)
    cp.set(run_section, "relative_files", "True")
    omit_elements = cp[run_section].get("omit", "").split("\n") or ["\n"]
    if "pytest.pex/*" not in omit_elements:
        omit_elements.append("pytest.pex/*")
    cp.set(run_section, "omit", "\n".join(omit_elements))
    stream = StringIO()
    cp.write(stream)
    return FileContent(fc.path, stream.getvalue().encode())


def get_branch_value_from_config(fc: FileContent) -> bool:
    # Note that coverage's default value for the branch setting is False, which we mirror here.
    if PurePath(fc.path).suffix == ".toml":
        all_config = _parse_toml_config(fc)
        return bool(
            all_config.get("tool", {}).get("coverage", {}).get("run", {}).get("branch", False)
        )

    cp = _parse_ini_config(fc)
    run_section = "coverage:run" if fc.path in ("tox.ini", "setup.cfg") else "run"
    if not cp.has_section(run_section):
        return False
    return cp.getboolean(run_section, "branch", fallback=False)


@rule
async def create_or_update_coverage_config(coverage: CoverageSubsystem) -> CoverageConfig:
    config_files = await Get(ConfigFiles, ConfigFilesRequest, coverage.config_request)
    if config_files.snapshot.files:
        digest_contents = await Get(DigestContents, Digest, config_files.snapshot.digest)
        file_content = _update_config(digest_contents[0])
    else:
        cp = configparser.ConfigParser()
        cp.add_section("run")
        cp.set("run", "relative_files", "True")
        cp.set("run", "omit", "\npytest.pex/*")
        stream = StringIO()
        cp.write(stream)
        file_content = FileContent(".coveragerc", stream.getvalue().encode())
    digest = await Get(Digest, CreateDigest([file_content]))
    return CoverageConfig(digest, file_content.path)


@dataclass(frozen=True)
class CoverageSetup:
    pex: VenvPex


@rule
async def setup_coverage(coverage: CoverageSubsystem) -> CoverageSetup:
    pex = await Get(VenvPex, PexRequest, coverage.to_pex_request())
    return CoverageSetup(pex)


@dataclass(frozen=True)
class MergedCoverageData:
    coverage_data: Digest
    addresses: tuple[Address, ...]


@rule(desc="Merge Pytest coverage data", level=LogLevel.DEBUG)
async def merge_coverage_data(
    data_collection: PytestCoverageDataCollection,
    coverage_setup: CoverageSetup,
    coverage_config: CoverageConfig,
    coverage: CoverageSubsystem,
    source_roots: AllSourceRoots,
) -> MergedCoverageData:
    if len(data_collection) == 1 and not coverage.global_report:
        coverage_data = data_collection[0]
        return MergedCoverageData(coverage_data.digest, (coverage_data.address,))

    coverage_digest_gets = []
    coverage_data_file_paths = []
    addresses = []
    for data in data_collection:
        # We prefix each .coverage file with its corresponding address to avoid collisions.
        coverage_digest_gets.append(
            Get(Digest, AddPrefix(data.digest, prefix=data.address.path_safe_spec))
        )
        coverage_data_file_paths.append(f"{data.address.path_safe_spec}/.coverage")
        addresses.append(data.address)

    if coverage.global_report:
        # It's important to set the `branch` value in the empty base report to the value it will
        # have when running on real inputs, so that the reports are of the same type, and can be
        # merged successfully. Otherwise we may get "Can't combine arc data with line data" errors.
        # See https://github.com/pantsbuild/pants/issues/14542 .
        config_contents = await Get(DigestContents, Digest, coverage_config.digest)
        branch = get_branch_value_from_config(config_contents[0]) if config_contents else False
        global_coverage_base_dir = PurePath("__global_coverage__")
        global_coverage_config_path = global_coverage_base_dir / "pyproject.toml"
        global_coverage_config_content = toml.dumps(
            {