# Copyright 2023 Pex project contributors.
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import

from pex.typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Any, Dict


def patch():
    # type: () -> None

    from pip._vendor.packaging import markers

    from pex.exceptions import production_assert
    from pex.pip.foreign_platform import EvaluationEnvironment, PatchContext

    evaluation_environment = PatchContext.load_evaluation_environment()

    def _get_env(
        environment,  # type: Dict[Any, Any]
        name,  # type: Any
    ):
        # type: (...) -> Any
        production_assert(
            isinstance(environment, EvaluationEnvironment),
            "Expected environment to come from the {function} function, "
            "which we patch to return {expected_type}, but was {actual_type}",
            function=markers.default_environment,
            expected_type=EvaluationEnvironment,
            actual_type=type(environment),
        )
        return environment[name]

    # Works with all Pip vendored packaging distributions.
    markers.default_environment = evaluation_environment.default
    # Covers Pip<24.1 vendored packaging.
    markers._get_env = _get_env

    original_eval_op = markers._eval_op

    def _eval_op(
        lhs,  # type: Any
        op,  # type: Any
        rhs,  # type: Any
        **kwargs  # type: Any
    ):
        # type: (...) -> Any
        evaluation_environment.raise_if_missing(lhs)
        evaluation_environment.raise_if_missing(rhs)
        return original_eval_op(lhs, op, rhs, **kwargs)

    markers._eval_op = _eval_op

    # Packaging 24.1+ unconditionally patches the env dict `python_full_version` value to work
    # around Python dev release versions with a `+` trailer not being valid PEP-440 versions. That
    # check will needlessly blow up for some of our AbbreviatedPlatform target environments; so we
    # short-circuit here.
    if "python_full_version" not in evaluation_environment:
        markers._repair_python_full_version = lambda env: env
