#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import annotations

import json

import pytest

from airflow.sdk._shared.module_loading import import_string
from airflow.sdk.configuration import conf, retrieve_configuration_description
from airflow.sdk.execution_time.coordinator import (
    BaseCoordinator,
    CoordinatorManager,
    _PythonCoordinator,
    get_coordinator_manager,
    reset_coordinator_manager,
)


class _CoordinatorA(BaseCoordinator):
    def __init__(self, *, label: str = "a"):
        self.label = label


class _CoordinatorB(BaseCoordinator):
    pass


class _ExplodingCoordinator(BaseCoordinator):
    def __init__(self):
        raise RuntimeError("This coordinator must not be instantiated")


@pytest.fixture
def sdk_config(monkeypatch):
    """Set the ``[sdk]`` env vars consumed by :meth:`CoordinatorManager.from_config`.

    :return: Callable ``apply(*, coordinators=None, queue_to_coordinator=None)`` --
        each argument is the raw JSON string for the matching env var, or ``None``
        to unset it. The conf cache is invalidated after each call (and again on
        teardown) so ``from_config()`` re-reads the values just set.
    """
    from airflow.sdk.configuration import conf

    def _apply(*, coordinators: str | None = None, queue_to_coordinator: str | None = None) -> None:
        if coordinators is None:
            monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False)
        else:
            monkeypatch.setenv("AIRFLOW__SDK__COORDINATORS", coordinators)
        if queue_to_coordinator is None:
            monkeypatch.delenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", raising=False)
        else:
            monkeypatch.setenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", queue_to_coordinator)
        conf.invalidate_cache()

    yield _apply
    conf.invalidate_cache()


class TestCoordinatorManager:
    @pytest.fixture(autouse=True)
    def _reset_cache(self):
        reset_coordinator_manager()
        yield
        reset_coordinator_manager()

    def test_from_config_loads_specs_and_resolves_instances(self, sdk_config):
        sdk_config(
            coordinators=json.dumps(
                {
                    "alpha": {
                        "classpath": f"{_CoordinatorA.__module__}._CoordinatorA",
                        "kwargs": {"label": "alpha-label"},
                    },
                    "beta": {"classpath": f"{_CoordinatorB.__module__}._CoordinatorB", "kwargs": {}},
                }
            ),
            queue_to_coordinator=json.dumps({"queue-a": "alpha"}),
        )
        manager = CoordinatorManager.from_config()
        assert manager._queue_to_coordinator == {"queue-a": "alpha"}
        assert manager._created_coordinators == {}

        coordinator_for_queue_a = manager.for_queue("queue-a")
        assert isinstance(coordinator_for_queue_a, _CoordinatorA)
        assert manager.for_queue("queue-a") is coordinator_for_queue_a, "instance should be cached"
        assert manager._created_coordinators == {"alpha": coordinator_for_queue_a}

        coordinator_for_queue_missing = manager.for_queue("queue-1")
        assert isinstance(coordinator_for_queue_missing, _PythonCoordinator)
        assert manager.for_queue("queue-1") is coordinator_for_queue_missing
        assert manager._created_coordinators == {"alpha": coordinator_for_queue_a}

    def test_from_config_empty(self, monkeypatch):
        monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False)
        monkeypatch.delenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", raising=False)
        conf.invalidate_cache()

        manager = CoordinatorManager.from_config()
        assert manager._coordinator_specs == {}
        assert manager._queue_to_coordinator == {}

    def test_get_coordinator_manager_is_cached(self, monkeypatch):
        monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False)

        from airflow.sdk.configuration import conf

        conf.invalidate_cache()

        m1 = get_coordinator_manager()
        m2 = get_coordinator_manager()
        assert m1 is m2

    @pytest.mark.parametrize(
        ("queue", "expected"),
        [
            pytest.param(
                "queue-java",
                {"pod_template_file": "/opt/airflow/pod_templates/java.yaml"},
                id="mapped-with-extra",
            ),
            pytest.param("queue-go", None, id="mapped-without-extra"),
            pytest.param("queue-unmapped", None, id="unmapped-queue"),
        ],
    )
    def test_extra_for_queue(self, sdk_config, queue, expected):
        sdk_config(
            coordinators=json.dumps(
                {
                    "java": {
                        "classpath": f"{_CoordinatorA.__module__}._CoordinatorA",
                        "extra": {"pod_template_file": "/opt/airflow/pod_templates/java.yaml"},
                    },
                    "go": {"classpath": f"{_CoordinatorB.__module__}._CoordinatorB"},
                }
            ),
            queue_to_coordinator=json.dumps({"queue-java": "java", "queue-go": "go"}),
        )
        manager = CoordinatorManager.from_config()
        # Resolving the extra must not instantiate the coordinator.
        assert manager.extra_for_queue(queue) == expected
        assert manager._created_coordinators == {}

    def test_extra_not_forwarded_to_constructor(self, sdk_config):
        """``extra`` is kept separate from ``kwargs`` and never reaches the coordinator constructor."""
        sdk_config(
            coordinators=json.dumps(
                {
                    "java": {
                        "classpath": f"{_CoordinatorA.__module__}._CoordinatorA",
                        "kwargs": {"label": "java-label"},
                        "extra": {"pod_template_file": "/opt/airflow/pod_templates/java.yaml"},
                    },
                }
            ),
            queue_to_coordinator=json.dumps({"queue-java": "java"}),
        )
        manager = CoordinatorManager.from_config()
        # _CoordinatorA only accepts ``label``; construction would raise TypeError
        # if ``extra`` were passed through.
        coordinator = manager.for_queue("queue-java")
        assert isinstance(coordinator, _CoordinatorA)
        assert coordinator.label == "java-label"
        # The extra is still readable from the spec without instantiation cost.
        assert manager.extra_for_queue("queue-java") == {
            "pod_template_file": "/opt/airflow/pod_templates/java.yaml"
        }

    def test_extra_for_queue_does_not_instantiate_coordinator(self, sdk_config):
        """Reading ``extra`` reads only the spec; a failing constructor must never run."""
        sdk_config(
            coordinators=json.dumps(
                {
                    "boom": {
                        "classpath": f"{_ExplodingCoordinator.__module__}._ExplodingCoordinator",
                        "extra": {"pod_template_file": "/opt/airflow/pod_templates/boom.yaml"},
                    },
                }
            ),
            queue_to_coordinator=json.dumps({"queue-boom": "boom"}),
        )
        manager = CoordinatorManager.from_config()
        assert manager.extra_for_queue("queue-boom") == {
            "pod_template_file": "/opt/airflow/pod_templates/boom.yaml"
        }
        assert manager._created_coordinators == {}


class TestConfigYamlCoordinatorsExample:
    """Guard the ``[sdk] coordinators`` example in ``config.yml`` against drift.

    Nothing else exercises the example, so a broken one (e.g. dropping the
    required ``jars_root`` kwarg) can ship unnoticed. Loading it through
    CoordinatorManager and constructing every entry keeps the example honest.
    """

    def test_every_example_coordinator_constructs(self, sdk_config):
        description = retrieve_configuration_description()
        coordinators_example = description["sdk"]["options"]["coordinators"]["example"]
        specs = json.loads(coordinators_example)
        assert specs, "config.yml [sdk] coordinators example must not be empty"

        # The example's own queue_to_coordinator illustrates different keys, so
        # route every coordinator through a synthetic queue to construct each one.
        queue_to_coordinator = {f"queue-{key}": key for key in specs}
        sdk_config(
            coordinators=coordinators_example,
            queue_to_coordinator=json.dumps(queue_to_coordinator),
        )
        manager = CoordinatorManager.from_config()
        assert set(manager._coordinator_specs) == set(specs)

        for queue, key in queue_to_coordinator.items():
            coordinator = manager.for_queue(queue)
            assert isinstance(coordinator, import_string(specs[key]["classpath"]))
