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

from __future__ import annotations

import dataclasses
import logging
from collections import defaultdict
from dataclasses import dataclass
from typing import Iterable

from packaging.utils import canonicalize_name as canonicalize_project_name

from pants.backend.python.pip_requirement import PipRequirement
from pants.backend.python.subsystems.setup import PythonSetup
from pants.backend.python.target_types import (
    MainSpecification,
    PexLayout,
    PythonRequirementsField,
    PythonResolveField,
    parse_requirements_file,
)
from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
from pants.backend.python.util_rules.local_dists import LocalDistsPex, LocalDistsPexRequest
from pants.backend.python.util_rules.local_dists import rules as local_dists_rules
from pants.backend.python.util_rules.pex import (
    CompletePlatforms,
    OptionalPex,
    OptionalPexRequest,
    PexPlatforms,
    PexRequest,
)
from pants.backend.python.util_rules.pex import rules as pex_rules
from pants.backend.python.util_rules.pex_requirements import Lockfile, PexRequirements
from pants.backend.python.util_rules.python_sources import (
    PythonSourceFiles,
    PythonSourceFilesRequest,
    StrippedPythonSourceFiles,
)
from pants.backend.python.util_rules.python_sources import rules as python_sources_rules
from pants.engine.addresses import Address, Addresses
from pants.engine.collection import DeduplicatedCollection
from pants.engine.fs import Digest, DigestContents, GlobMatchErrorBehavior, MergeDigests, PathGlobs
from pants.engine.rules import Get, MultiGet, collect_rules, rule
from pants.engine.target import Target, TransitiveTargets, TransitiveTargetsRequest
from pants.util.docutil import doc_url
from pants.util.logging import LogLevel
from pants.util.meta import frozen_after_init
from pants.util.strutil import bullet_list, path_safe

logger = logging.getLogger(__name__)


@frozen_after_init
@dataclass(unsafe_hash=True)
class PexFromTargetsRequest:
    addresses: Addresses
    output_filename: str
    internal_only: bool
    layout: PexLayout | None
    main: MainSpecification | None
    platforms: PexPlatforms
    complete_platforms: CompletePlatforms
    additional_args: tuple[str, ...]
    additional_lockfile_args: tuple[str, ...]
    include_source_files: bool
    include_requirements: bool
    include_local_dists: bool
    additional_sources: Digest | None
    additional_inputs: Digest | None
    hardcoded_interpreter_constraints: InterpreterConstraints | None
    # This field doesn't participate in comparison (and therefore hashing), as it doesn't affect
    # the result.
    description: str | None = dataclasses.field(compare=False)

    def __init__(
        self,
        addresses: Iterable[Address],
        *,
        output_filename: str,
        internal_only: bool,
        layout: PexLayout | None = None,
        main: MainSpecification | None = None,
        platforms: PexPlatforms = PexPlatforms(),
        complete_platforms: CompletePlatforms = CompletePlatforms(),
        additional_args: Iterable[str] = (),
        additional_lockfile_args: Iterable[str] = (),
        include_source_files: bool = True,
        include_requirements: bool = True,
        include_local_dists: bool = False,
        additional_sources: Digest | None = None,
        additional_inputs: Digest | None = None,
        hardcoded_interpreter_constraints: InterpreterConstraints | None = None,
        description: str | None = None,
    ) -> None:
        """Request to create a Pex from the transitive closure of the given addresses.

        :param addresses: The addresses to use for determining what is included in the Pex. The
            transitive closure of these addresses will be used; you only need to specify the roots.
        :param output_filename: The name of the built Pex file, which typically should end in
            `.pex`.
        :param internal_only: Whether we ever materialize the Pex and distribute it directly
            to end users, such as with the `binary` goal. Typically, instead, the user never
            directly uses the Pex, e.g. with `lint` and `test`. If True, we will use a Pex setting
            that results in faster build time but compatibility with fewer interpreters at runtime.
        :param layout: The filesystem layout to create the PEX with.
        :param main: The main for the built Pex, equivalent to Pex's `-e` or `-c` flag. If
            left off, the Pex will open up as a REPL.
        :param platforms: Which platforms should be supported. Setting this value will cause
            interpreter constraints to not be used because platforms already constrain the valid
            Python versions, e.g. by including `cp36m` in the platform string.
        :param additional_args: Any additional Pex flags.
        :param additional_lockfile_args: Any additional Pex flags that should be used with the
            lockfile.pex. Many Pex args like `--emit-warnings` do not impact the lockfile, and
            setting them would reduce reuse with other call sites. Generally, these should only be
            flags that impact lockfile resolution like `--manylinux`.
        :param include_source_files: Whether to include source files in the built Pex or not.
            Setting this to `False` and loading the source files by instead populating the chroot
            and setting the environment variable `PEX_EXTRA_SYS_PATH` will result in substantially
            fewer rebuilds of the Pex.
        :param include_requirements: Whether to resolve requirements and include them in the Pex.
        :param include_local_dists: Whether to build local dists and include them in the built pex.
        :param additional_sources: Any additional source files to include in the built Pex.
        :param additional_inputs: Any inputs that are not source files and should not be included
            directly in the Pex, but should be present in the environment when building the Pex.
        :param hardcoded_interpreter_constraints: Use these constraints rather than resolving the
            constraints from the input.
        :param description: A human-readable description to render in the dynamic UI when building
            the Pex.
        """
        self.addresses = Addresses(addresses)
        self.output_filename = output_filename
        self.internal_only = internal_only
        self.layout = layout
        self.main = main
        self.platforms = platforms
        self.complete_platforms = complete_platforms
        self.additional_args = tuple(additional_args)
        self.additional_lockfile_args = tuple(additional_lockfile_args)
        self.include_source_files = include_source_files
        self.include_requirements = include_requirements
        self.include_local_dists = include_local_dists
        self.additional_sources = additional_sources
        self.additional_inputs = additional_inputs
        self.hardcoded_interpreter_constraints = hardcoded_interpreter_constraints
        self.description = description

    def to_interpreter_constraints_request(self) -> InterpreterConstraintsRequest:
        return InterpreterConstraintsRequest(
            addresses=self.addresses,
            hardcoded_interpreter_constraints=self.hardcoded_interpreter_constraints,
        )


@frozen_after_init
@dataclass(unsafe_hash=True)
class InterpreterConstraintsRequest:
    addresses: Addresses
    hardcoded_interpreter_constraints: InterpreterConstraints | None

    def __init__(
        self,
        addresses: Iterable[Address],
        *,
        hardcoded_interpreter_constraints: InterpreterConstraints | None = None,
    ) -> None:
        self.addresses = Addresses(addresses)
        self.hardcoded_interpreter_constraints = hardcoded_interpreter_constraints


@rule
async def interpreter_constraints_for_targets(
    request: InterpreterConstraintsRequest, python_setup: PythonSetup
) -> InterpreterConstraints:
    if request.hardcoded_interpreter_constraints:
        return request.hardcoded_interpreter_constraints

    transitive_targets = await Get(TransitiveTargets, TransitiveTargetsRequest(request.addresses))
    calculated_constraints = InterpreterConstraints.create_from_targets(
        transitive_targets.closure, python_setup
    )
    # If there are no targets, we fall back to the global constraints. This is relevant,
    # for example, when running `./pants repl` with no specs or only on targets without
    # `interpreter_constraints` (e.g. `python_requirement`).
    interpreter_constraints = calculated_constraints or InterpreterConstraints(
        python_setup.interpreter_constraints
    )
    return interpreter_constraints


@dataclass(frozen=True)
class ChosenPythonResolve:
    name: str
    lockfile_path: str


@dataclass(frozen=True)
class ChosenPythonResolveRequest:
    addresses: Addresses


# Note: Inspired by `coursier_fetch.py`.
class NoCompatibleResolveException(Exception):
    """No compatible resolve could be found for a set of targets."""

    def __init__(
        self,
        python_setup: PythonSetup,
        msg_prefix: str,
        relevant_targets: Iterable[Target],
        msg_suffix: str = "",
    ) -> None:
        resolves_to_addresses = defaultdict(list)
        for tgt in relevant_targets:
            if tgt.has_field(PythonResolveField):
                resolve = tgt[PythonResolveField].normalized_value(python_setup)
                resolves_to_addresses[resolve].append(tgt.address.spec)

        formatted_resolve_lists = "\n\n".join(
            f"{resolve}:\n{bullet_list(sorted(addresses))}"
            for resolve, addresses in sorted(resolves_to_addresses.items())
        )
        super().__init__(
            f"{msg_prefix}:\n\n"
            f"{formatted_resolve_lists}\n\n"
            "Targets which will be used together must all have the same resolve (from the "
            f"[resolve]({doc_url('reference-python_test#coderesolvecode')}) "
            "field) in common." + (f"\n\n{msg_suffix}" if msg_suffix else "")
        )


@rule
async def choose_python_resolve(
    request: ChosenPythonResolveRequest, python_setup: PythonSetup
) -> ChosenPythonResolve:
    transitive_targets = await Get(TransitiveTargets, TransitiveTargetsRequest(request.addresses))

    # First, choose the resolve by inspecting the root targets.
    root_resolves = {
        root[PythonResolveField].normalized_value(python_setup)
        for root in transitive_targets.roots
        if root.has_field(PythonResolveField)
    }
    if not root_resolves:
        # If there are no relevant targets, we fall back to the default resolve. This is relevant,
        # for example, when running `./pants repl` with no specs or only on non-Python targets.
        return ChosenPythonResolve(
            name=python_setup.default_resolve,
            lockfile_path=python_setup.resolves[python_setup.default_resolve],
        )

    if len(root_resolves) > 1:
        raise NoCompatibleResolveException(
            python_setup,
            "The input targets did not have a resolve in common",
            transitive_targets.roots,
        )

    chosen_resolve = next(iter(root_resolves))

    # Then, validate that all transitive deps are compatible.
    for tgt in transitive_targets.dependencies:
        if (
            tgt.has_field(PythonResolveField)
            and tgt[PythonResolveField].normalized_value(python_setup) != chosen_resolve
        ):
            plural = ("s", "their") if len(transitive_targets.roots) > 1 else ("", "its")
            raise NoCompatibleResolveException(
                python_setup,
                (
                    f"The resolve chosen for the root target{plural[0]} was {chosen_resolve}, but "
                    f"some of {plural[1]} dependencies are not compatible with that resolve"
                ),
                transitive_targets.closure,
            )

    return ChosenPythonResolve(
        name=chosen_resolve, lockfile_path=python_setup.resolves[chosen_resolve]
    )


class GlobalRequirementConstraints(DeduplicatedCollection[PipRequirement]):
    """Global constraints specified by the `[python].requirement_constraints` setting, if any."""


@rule
async def determine_global_requirement_constraints(
    python_setup: PythonSetup,
) -> GlobalRequirementConstraints:
    if not python_setup.requirement_constraints:
        return GlobalRequirementConstraints()

    constraints_file_contents = await Get(
        DigestContents,
        PathGlobs(
            [python_setup.requirement_constraints],