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

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from textwrap import dedent
from typing import Any

import pytest

from pants.build_graph.address import Address, ResolveError
from pants.core.environments import rules as environments_rules
from pants.core.environments.rules import (
    AllEnvironmentTargets,
    AmbiguousEnvironmentError,
    ChosenLocalEnvironmentName,
    EnvironmentName,
    EnvironmentNameRequest,
    NoFallbackEnvironmentError,
    SingleEnvironmentNameRequest,
    UnrecognizedEnvironmentError,
    extract_process_config_from_environment,
    resolve_environment_name,
    resolve_environment_sensitive_option,
)
from pants.core.environments.subsystems import EnvironmentsSubsystem
from pants.core.environments.target_types import (
    CompatiblePlatformsField,
    DockerEnvironmentTarget,
    DockerImageField,
    DockerPlatformField,
    EnvironmentField,
    EnvironmentTarget,
    FallbackEnvironmentField,
    LocalEnvironmentTarget,
    LocalWorkspaceEnvironmentTarget,
    RemoteEnvironmentCacheBinaryDiscovery,
    RemoteEnvironmentTarget,
    RemoteExtraPlatformPropertiesField,
)
from pants.engine.environment import LOCAL_ENVIRONMENT_MATCHER, ChosenLocalWorkspaceEnvironmentName
from pants.engine.internals.docker import DockerResolveImageResult
from pants.engine.platform import Platform
from pants.engine.process import ProcessCacheScope
from pants.engine.target import FieldSet, OptionalSingleSourceField, Target
from pants.option.global_options import GlobalOptions, KeepSandboxes
from pants.option.option_types import ShellStrListOption
from pants.option.subsystem import Subsystem
from pants.testutil.option_util import create_subsystem
from pants.testutil.rule_runner import QueryRule, RuleRunner, engine_error, run_rule_with_mocks


@pytest.fixture
def rule_runner() -> RuleRunner:
    return RuleRunner(
        rules=[
            *environments_rules.rules(),
            QueryRule(AllEnvironmentTargets, []),
            QueryRule(EnvironmentTarget, [EnvironmentName]),
            QueryRule(EnvironmentName, [EnvironmentNameRequest]),
            QueryRule(EnvironmentName, [SingleEnvironmentNameRequest]),
            QueryRule(ChosenLocalEnvironmentName, []),
            QueryRule(ChosenLocalWorkspaceEnvironmentName, []),
        ],
        target_types=[
            LocalEnvironmentTarget,
            LocalWorkspaceEnvironmentTarget,
            DockerEnvironmentTarget,
            RemoteEnvironmentTarget,
        ],
        inherent_environment=None,
    )


def test_extract_process_config_from_environment() -> None:
    def assert_config(
        *,
        envs_enabled: bool = True,
        env_tgt: (
            LocalEnvironmentTarget
            | LocalWorkspaceEnvironmentTarget
            | RemoteEnvironmentTarget
            | DockerEnvironmentTarget
            | None
        ),
        enable_remote_execution: bool,
        expected_remote_execution: bool,
        expected_docker_image: str | None,
        expected_remote_execution_extra_platform_properties: list[tuple[str, str]] | None = None,
    ) -> None:
        global_options = create_subsystem(
            GlobalOptions,
            remote_execution=enable_remote_execution,
            remote_execution_extra_platform_properties=["global_k=v"],
        )
        name = "name"
        env_subsystem = create_subsystem(
            EnvironmentsSubsystem,
            names={name: "addr"} if envs_enabled else {},
        )
        keep_sandboxes = KeepSandboxes.on_failure
        result = run_rule_with_mocks(
            extract_process_config_from_environment,
            rule_args=[
                EnvironmentTarget(name, env_tgt),
                Platform.linux_arm64,
                global_options,
                keep_sandboxes,
                env_subsystem,
            ],
            mock_calls={
                "pants.engine.intrinsics.docker_resolve_image": lambda req: DockerResolveImageResult(
                    "sha256:abc123"
                ),
            },
        )
        assert result.platform == Platform.linux_arm64.value
        assert result.remote_execution is expected_remote_execution
        if expected_docker_image is not None:
            # TODO(17104): Account for any prepended image name here once implemented.
            assert result.docker_image == "sha256:abc123"
        else:
            assert result.docker_image is None
        assert result.remote_execution_extra_platform_properties == (
            expected_remote_execution_extra_platform_properties or []
        )

    assert_config(
        env_tgt=None,
        enable_remote_execution=False,
        expected_remote_execution=False,
        expected_docker_image=None,
    )
    assert_config(
        env_tgt=None,
        envs_enabled=False,
        enable_remote_execution=True,
        expected_remote_execution=True,
        expected_docker_image=None,
        expected_remote_execution_extra_platform_properties=[("global_k", "v")],
    )
    assert_config(
        env_tgt=None,
        envs_enabled=True,
        enable_remote_execution=True,
        expected_remote_execution=False,
        expected_docker_image=None,
    )

    for re in (False, True):
        assert_config(
            env_tgt=LocalEnvironmentTarget({}, Address("dir")),
            enable_remote_execution=re,
            expected_remote_execution=False,
            expected_docker_image=None,
        )

    for re in (False, True):
        assert_config(
            env_tgt=DockerEnvironmentTarget(
                {
                    DockerImageField.alias: "my_img",
                },
                Address("dir"),
            ),
            enable_remote_execution=re,
            expected_remote_execution=False,
            expected_docker_image="my_img",
        )

    for re in (False, True):
        assert_config(
            env_tgt=RemoteEnvironmentTarget({}, Address("dir")),
            enable_remote_execution=re,
            expected_remote_execution=True,
            expected_docker_image=None,
            # The global option is ignored.
            expected_remote_execution_extra_platform_properties=[],
        )
    assert_config(
        env_tgt=RemoteEnvironmentTarget(
            {RemoteExtraPlatformPropertiesField.alias: ["field_k=v"]}, Address("dir")
        ),
        enable_remote_execution=re,
        expected_remote_execution=True,
        expected_docker_image=None,
        expected_remote_execution_extra_platform_properties=[("field_k", "v")],
    )


def test_all_environments(rule_runner: RuleRunner) -> None:
    rule_runner.write_files(
        {
            "BUILD": dedent(
                """\
                local_environment(name='e1')
                local_environment(name='e2')
                local_environment(name='no-name')
                docker_environment(name='docker', image="centos6:latest")
                """
            )
        }
    )
    rule_runner.set_options(
        ["--environments-preview-names={'e1': '//:e1', 'e2': '//:e2', 'docker': '//:docker'}"]
    )
    result = rule_runner.request(AllEnvironmentTargets, [])
    assert {i: j.address for (i, j) in result.items()} == {
        "e1": Address("", target_name="e1"),
        "e2": Address("", target_name="e2"),
        "docker": Address("", target_name="docker"),
    }


def test_choose_local_environment(rule_runner: RuleRunner) -> None:
    rule_runner.write_files(
        {
            "BUILD": dedent(
                """\
                local_environment(name='e1')
                local_environment(name='e2')
                local_environment(name='not-compatible', compatible_platforms=[])
                docker_environment(name='docker', docker_image="centos6:latest")
                """
            )
        }
    )

    def get_env() -> EnvironmentTarget:
        name = rule_runner.request(ChosenLocalEnvironmentName, [])
        return rule_runner.request(EnvironmentTarget, [name.val])

    # If `--names` is not set, do not choose an environment.
    assert get_env().val is None

    rule_runner.set_options(["--environments-preview-names={'e': '//:e1'}"])
    assert get_env().val.address == Address("", target_name="e1")  # type: ignore[union-attr]

    # If `--names` is set, but no compatible platforms, do not choose an environment.
    rule_runner.set_options(["--environments-preview-names={'e': '//:not-compatible'}"])
    assert get_env().val is None

    # Error if >1 compatible targets.
    rule_runner.set_options(["--environments-preview-names={'e1': '//:e1', 'e2': '//:e2'}"])
    with engine_error(AmbiguousEnvironmentError):
        get_env()


def test_resolve_environment_name(rule_runner: RuleRunner) -> None:
    rule_runner.write_files(
        {
            "BUILD": dedent(
                """\
                local_environment(name='local')
                local_environment(
                    name='local-fallback', compatible_platforms=[], fallback_environment='local'
                )
                docker_environment(name='docker', image="centos6:latest")
                remote_environment(name='remote-no-fallback')
                remote_environment(name='remote-fallback', fallback_environment="docker")
                remote_environment(name='remote-bad-fallback', fallback_environment="fake")
                """
            )
        }
    )

    def get_name(v: str) -> EnvironmentName:
        return rule_runner.request(
            EnvironmentName, [EnvironmentNameRequest(v, description_of_origin="foo")]
        )

    # If `--names` is not set, and the local matcher is used, do not choose an environment.
    assert get_name(LOCAL_ENVIRONMENT_MATCHER).val is None
    # Else, error for unrecognized names.
    with engine_error(UnrecognizedEnvironmentError):
        get_name("local")

    env_names_arg = (
        "--environments-preview-names={"
        + "'local': '//:local', "
        + "'local-fallback': '//:local-fallback', "
        + "'docker': '//:docker', "
        + "'remote-no-fallback': '//:remote-no-fallback', "
        + "'remote-fallback': '//:remote-fallback',"
        "'remote-bad-fallback': '//:remote-bad-fallback'}"
    )
    rule_runner.set_options([env_names_arg, "--remote-execution"])
    assert get_name(LOCAL_ENVIRONMENT_MATCHER).val == "local"
    for name in ("local", "docker", "remote-no-fallback", "remote-fallback"):
        assert get_name(name).val == name
    with engine_error(UnrecognizedEnvironmentError):
        get_name("fake")

    assert get_name("local-fallback").val == "local"

    rule_runner.set_options([env_names_arg, "--no-remote-execution"])
    assert get_name("remote-fallback").val == "docker"
    with engine_error(NoFallbackEnvironmentError):
        get_name("remote-no-fallback")
    with engine_error(UnrecognizedEnvironmentError):
        get_name("remote-bad-fallback")


def test_resolve_environment_names(rule_runner: RuleRunner) -> None:
    rule_runner.write_files(
        {
            "BUILD": dedent(
                """\
                local_environment(name='local1')
                local_environment(name='local2')
                """
            )
        }
    )

    def get_names(v: Sequence[str]) -> EnvironmentName:
        return rule_runner.request(
            EnvironmentName, [SingleEnvironmentNameRequest(tuple(v), description_of_origin="foo")]
        )

    # If `--names` is not set, and the local matcher is used, do not choose an environment.
    assert get_names([LOCAL_ENVIRONMENT_MATCHER]).val is None

    env_names_arg = (
        "--environments-preview-names={" + "'local1': '//:local1', " + "'local2': '//:local2'}"
    )
    rule_runner.set_options([env_names_arg])

    assert get_names(["local1", "local1"]).val == "local1"

    with engine_error(contains="Needed 1 unique environment, but foo contained 2:"):
        _ = get_names(["local1", "local2"])


def test_resolve_environment_name_local_and_docker_fallbacks(monkeypatch) -> None:
    # We can't monkeypatch the Platform with RuleRunner, so instead use run_rule_with_mocks.
    def get_env_name(
        env_tgt: Target, platform: Platform, *, docker_execution: bool = True
    ) -> str | None:
        def mock_determine_local_environment() -> ChosenLocalEnvironmentName:
            return ChosenLocalEnvironmentName(EnvironmentName(None))

        def mock_determine_local_workspace_environment() -> ChosenLocalWorkspaceEnvironmentName:
            return ChosenLocalWorkspaceEnvironmentName(EnvironmentName(None))

        def mock_get_target_for_environment_name(name: EnvironmentName) -> EnvironmentTarget:
            return EnvironmentTarget(name.val or "__local__", env_tgt)

        monkeypatch.setattr(Platform, "create_for_localhost", lambda: platform)
        result = run_rule_with_mocks(
            resolve_environment_name,
            rule_args=[
                EnvironmentNameRequest("env", description_of_origin="foo"),
                create_subsystem(EnvironmentsSubsystem, names={"env": "", "fallback": ""}),
                create_subsystem(
                    GlobalOptions, remote_execution=False, docker_execution=docker_execution
                ),
            ],
            mock_calls={
                "pants.core.environments.rules.determine_local_environment": mock_determine_local_environment,
                "pants.core.environments.rules.determine_local_workspace_environment": mock_determine_local_workspace_environment,
                "pants.core.environments.rules.get_target_for_environment_name": mock_get_target_for_environment_name,
                "pants.core.environments.rules.resolve_environment_name": lambda req: EnvironmentName(
                    req.raw_value
                ),
            },
        ).val
        return result

    def create_local_tgt(
        *, compatible_platforms: list[Platform] | None = None, fallback: bool = False
    ) -> LocalEnvironmentTarget:
        return LocalEnvironmentTarget(
            {
                CompatiblePlatformsField.alias: (
                    [plat.value for plat in compatible_platforms] if compatible_platforms else None
                ),
                FallbackEnvironmentField.alias: "fallback" if fallback else None,
            },
            Address("envs"),
        )

    assert get_env_name(create_local_tgt(), Platform.linux_arm64) == "env"
    assert (
        get_env_name(
            create_local_tgt(compatible_platforms=[Platform.linux_x86_64], fallback=True),
            Platform.linux_arm64,
        )
        == "fallback"
    )
    with pytest.raises(NoFallbackEnvironmentError):
        get_env_name(
            create_local_tgt(compatible_platforms=[Platform.linux_x86_64]),
            Platform.linux_arm64,
        )

    def create_docker_tgt(
        *, platform: Platform | None = None, fallback: bool = False
    ) -> DockerEnvironmentTarget:
        return DockerEnvironmentTarget(
            {
                DockerImageField.alias: "image",
                DockerPlatformField.alias: platform.value if platform else None,
                FallbackEnvironmentField.alias: "fallback" if fallback else None,
            },
            Address("envs"),
        )

    # If the Docker platform is not set, we default to the CPU arch. So regardless of localhost,
    # the Docker environment should be used.
    assert get_env_name(create_docker_tgt(), Platform.linux_arm64) == "env"

    # If Docker execution is disabled, though, fallback.
    assert (
        get_env_name(create_docker_tgt(fallback=True), Platform.linux_arm64, docker_execution=False)
        == "fallback"
    )
    with pytest.raises(NoFallbackEnvironmentError):
        get_env_name(create_docker_tgt(), Platform.linux_arm64, docker_execution=False)

    # The Docker env can be used if we're on macOS, or on Linux and the CPU arch matches.
    for plat in (Platform.macos_arm64, Platform.macos_x86_64, Platform.linux_x86_64):
        assert get_env_name(create_docker_tgt(platform=Platform.linux_x86_64), plat) == "env"
    for plat in (Platform.macos_arm64, Platform.macos_x86_64, Platform.linux_arm64):
        assert get_env_name(create_docker_tgt(platform=Platform.linux_arm64), plat) == "env"

    # But if on Linux and a different CPU arch is used, fallback.
    assert (
        get_env_name(
            create_docker_tgt(platform=Platform.linux_x86_64, fallback=True), Platform.linux_arm64
        )
        == "fallback"
    )
    assert (
        get_env_name(
            create_docker_tgt(platform=Platform.linux_arm64, fallback=True), Platform.linux_x86_64
        )
        == "fallback"
    )
    with pytest.raises(NoFallbackEnvironmentError):
        get_env_name(create_docker_tgt(platform=Platform.linux_x86_64), Platform.linux_arm64)
    with pytest.raises(NoFallbackEnvironmentError):
        get_env_name(create_docker_tgt(platform=Platform.linux_arm64), Platform.linux_x86_64)


def test_resolve_environment_tgt(rule_runner: RuleRunner) -> None:
    rule_runner.write_files({"BUILD": "local_environment(name='env')"})
    rule_runner.set_options(
        ["--environments-preview-names={'env': '//:env', 'bad-address': '//:fake'}"]
    )

    def get_tgt(v: str | None) -> EnvironmentTarget:
        return rule_runner.request(EnvironmentTarget, [EnvironmentName(v)])

    assert get_tgt(None).val is None

    # If the name is not defined, error.
    with engine_error(AssertionError):
        get_tgt("bad")

    assert get_tgt("env").val.address == Address("", target_name="env")  # type: ignore[union-attr]
    with engine_error(ResolveError):
        get_tgt("bad-address")


def test_environment_name_request_from_field_set() -> None:
    class EnvFieldSubclass(EnvironmentField):
        alias = "the_env_field"  # This intentionally uses a custom alias.

    class Tgt(Target):
        alias = "tgt"
        help = "foo"
        core_fields = (OptionalSingleSourceField, EnvFieldSubclass)

    @dataclass(frozen=True)
    class NoEnvFS(FieldSet):
        required_fields = (OptionalSingleSourceField,)

        source: OptionalSingleSourceField

    @dataclass(frozen=True)
    class EnvFS(FieldSet):
        required_fields = (OptionalSingleSourceField,)

        source: OptionalSingleSourceField
        the_env_field: EnvironmentField  # This intentionally uses an unusual attribute name.

    tgt = Tgt({EnvFieldSubclass.alias: "my_env"}, Address("dir"))
    assert EnvironmentNameRequest.from_field_set(NoEnvFS.create(tgt)) == EnvironmentNameRequest(
        LOCAL_ENVIRONMENT_MATCHER,
        description_of_origin="the `environment` field from the target dir:dir",
    )
    assert EnvironmentNameRequest.from_field_set(EnvFS.create(tgt)) == EnvironmentNameRequest(
        "my_env", description_of_origin="the `the_env_field` field from the target dir:dir"
    )


def test_executable_search_path_cache_scope() -> None:
    def assert_cache(
        tgt: Target | None, *, cache_failures: bool, expected: ProcessCacheScope
    ) -> None:
        assert (
            EnvironmentTarget("unused", tgt).executable_search_path_cache_scope(
                cache_failures=cache_failures
            )
            == expected
        )

    addr = Address("envs")

    for tgt in (
        None,
        LocalEnvironmentTarget({}, addr),
        RemoteEnvironmentTarget({RemoteEnvironmentCacheBinaryDiscovery.alias: False}, addr),
    ):
        assert_cache(tgt, cache_failures=False, expected=ProcessCacheScope.PER_RESTART_SUCCESSFUL)
        assert_cache(tgt, cache_failures=True, expected=ProcessCacheScope.PER_RESTART_ALWAYS)

    for tgt in (  # type: ignore[assignment]
        DockerEnvironmentTarget({DockerImageField.alias: "img"}, addr),
        RemoteEnvironmentTarget({RemoteEnvironmentCacheBinaryDiscovery.alias: True}, addr),
    ):
        assert_cache(tgt, cache_failures=False, expected=ProcessCacheScope.SUCCESSFUL)
        assert_cache(tgt, cache_failures=True, expected=ProcessCacheScope.ALWAYS)


# Test for regression in choosing local environments.
def test_find_chosen_local_and_experimental_workspace_environments(rule_runner: RuleRunner) -> None:
    rule_runner.write_files(
        {
            "BUILD": dedent(
                """\
            local_environment(name="local")
            experimental_workspace_environment(name="workspace")
            """
            )
        }
    )
    rule_runner.set_options(
        ["--environments-preview-names={'local': '//:local', 'workspace': '//:workspace'}"]
    )

    chosen_local_env = rule_runner.request(ChosenLocalEnvironmentName, [])
    assert chosen_local_env.val.val == "local"

    chosen_workspace_env = rule_runner.request(ChosenLocalWorkspaceEnvironmentName, [])
    assert chosen_workspace_env.val.val == "workspace"


def test_sequence_field_inheritance() -> None:
    """Test environment field inheritance for sequence fields."""

    class TestSubsystem(Subsystem):
        options_scope = "test-subsystem"
        help = "Test subsystem"

        class EnvironmentAware(Subsystem.EnvironmentAware):
            test_sequence_option = ShellStrListOption(help="Test sequence option")

    rule_runner = RuleRunner(
        target_types=[LocalEnvironmentTarget],
        rules=[
            *TestSubsystem.rules(),
            *environments_rules.rules(),
        ],
    )

    rule_runner.set_options(
        [
            "--test-subsystem-test-sequence-option=GLOBAL_VAR=global_value",
            "--test-subsystem-test-sequence-option=TEST_VAR=test_value",
        ]
    )

    def test_environment_resolution(
        runner: RuleRunner,
        env_name: str,
        env_definition: str,
        expected_result: Any,
        description: str,
    ):
        runner.write_files(
            {
                f"{env_name}/BUILD": dedent(f"""\
            local_environment(
                name="{env_name}",
                compatible_platforms=["linux_x86_64", "macos_arm64"],
                {env_definition}
            )
            """)
            }
        )

        env_target = runner.get_target(Address(env_name, target_name=env_name))

        test_sequence_field: Any = None
        for field_type in env_target.field_values:
            field = env_target.field_values[field_type]
            if hasattr(field, "alias") and field.alias == "test_subsystem_test_sequence_option":
                test_sequence_field = field
                break

        assert test_sequence_field is not None

        subsystem_class = test_sequence_field.subsystem
        mock_env_aware = subsystem_class()
        mock_env_aware.env_tgt = EnvironmentTarget(env_name, env_target)

        result = resolve_environment_sensitive_option(
            test_sequence_field.option_name,
            mock_env_aware,
        )

        assert result == expected_result, f"{description}: Expected {expected_result}, got {result}"

    test_environment_resolution(
        rule_runner, "env1", "", None, "Field not specified - should inherit global config"
    )

    test_environment_resolution(
        rule_runner,
        "env2",
        "test_subsystem_test_sequence_option=[],",
        (),
        "Global values can be overridden with empty",
    )

    test_environment_resolution(
        rule_runner,
        "env3",
        'test_subsystem_test_sequence_option=["ENV_VAR=env_value"],',
        ("ENV_VAR=env_value",),
        "Field with values - should override global config",
    )

    # Set global options to empty to test inheritance from empty global
    rule_runner.set_options([])

    test_environment_resolution(
        rule_runner,
        "env4",
        "",
        None,
        "Field not specified with empty global - should inherit empty global",
    )

    test_environment_resolution(
        rule_runner,
        "env5",
        'test_subsystem_test_sequence_option=["NEW_VAR=new_value"],',
        ("NEW_VAR=new_value",),
        "Empty global can be overridden with values",
    )
