from collections.abc import Iterator

import dagster as dg
import pytest
from dagster._core.events import DagsterEventType
from dagster._core.instance import DagsterInstance
from dagster._core.remote_representation.external import RemoteJob
from dagster._core.storage.dagster_run import DagsterRunStatus
from dagster._core.test_utils import create_run_for_test, environ
from dagster._core.workspace.context import WorkspaceRequestContext
from dagster._utils.merger import merge_dicts

from dagster_tests.api_tests.utils import get_bar_workspace


class TestQueuedRunCoordinator:
    """You can extend this class to easily run these set of tests on any custom run coordinator
    that subclasses the QueuedRunCoordinator. When extending, you simply need to override the
    `coordinator` fixture and return your implementation of `QueuedRunCoordinator`.

    For example:

    ```
    class TestMyRunCoordinator(TestQueuedRunCoordinator):
        @pytest.fixture(scope='function')
        def coordinator(self, instance):
            run_coordinator = MyRunCoordinator()
            run_coordinator.register_instance(instance)
            yield run_coordinator
    ```
    """

    @pytest.fixture
    def instance(self) -> Iterator[dg.DagsterInstance]:
        overrides = {
            "run_launcher": {"module": "dagster._core.test_utils", "class": "MockedRunLauncher"}
        }
        with dg.instance_for_test(overrides=overrides) as inst:
            yield inst

    @pytest.fixture
    def coordinator(self, instance: DagsterInstance) -> Iterator[dg.QueuedRunCoordinator]:
        run_coordinator = dg.QueuedRunCoordinator()
        run_coordinator.register_instance(instance)
        yield run_coordinator

    @pytest.fixture(name="workspace")
    def workspace_fixture(self, instance: DagsterInstance) -> Iterator[WorkspaceRequestContext]:
        with get_bar_workspace(instance) as workspace:
            yield workspace

    @pytest.fixture(name="remote_job")
    def remote_job_fixture(self, workspace: WorkspaceRequestContext) -> RemoteJob:
        location = workspace.get_code_location("bar_code_location")
        return location.get_repository("bar_repo").get_full_job("foo")

    def create_run_for_test(
        self, instance: DagsterInstance, remote_job: RemoteJob, **kwargs: object
    ) -> dg.DagsterRun:
        job_args = merge_dicts(
            {
                "job_name": "foo",
                "remote_job_origin": remote_job.get_remote_origin(),
                "job_code_origin": remote_job.get_python_origin(),
            },
            kwargs,
        )
        return create_run_for_test(instance, **job_args)

    def test_config(self):
        with environ({"MAX_RUNS": "10", "DEQUEUE_INTERVAL": "7"}):
            with dg.instance_for_test(
                overrides={
                    "run_coordinator": {
                        "module": "dagster._core.run_coordinator",
                        "class": "QueuedRunCoordinator",
                        "config": {
                            "max_concurrent_runs": {
                                "env": "MAX_RUNS",
                            },
                            "tag_concurrency_limits": [
                                {"key": "foo", "value": "bar", "limit": 3},
                                {"key": "backfill", "limit": 2},
                            ],
                            "dequeue_interval_seconds": {
                                "env": "DEQUEUE_INTERVAL",
                            },
                            "block_op_concurrency_limited_runs": {
                                "enabled": True,
                                "op_concurrency_slot_buffer": 1,
                            },
                        },
                    }
                }
            ) as _:
                pass

        with pytest.raises(dg.DagsterInvalidConfigError):
            with dg.instance_for_test(
                overrides={
                    "run_coordinator": {
                        "module": "dagster._core.run_coordinator",
                        "class": "QueuedRunCoordinator",
                        "config": {
                            "tag_concurrency_limits": [
                                {"key": "backfill"},
                            ],
                        },
                    }
                }
            ) as instance:
                print(instance.run_coordinator)  # noqa: T201

    def test_config_unique_value(self):
        with environ({"MAX_RUNS": "10", "DEQUEUE_INTERVAL": "7"}):
            with dg.instance_for_test(
                overrides={
                    "run_coordinator": {
                        "module": "dagster._core.run_coordinator",
                        "class": "QueuedRunCoordinator",
                        "config": {
                            "max_concurrent_runs": {
                                "env": "MAX_RUNS",
                            },
                            "tag_concurrency_limits": [
                                {
                                    "key": "foo",
                                    "value": {"applyLimitPerUniqueValue": True},
                                    "limit": 3,
                                },
                                {"key": "backfill", "limit": 2},
                            ],
                            "dequeue_interval_seconds": {
                                "env": "DEQUEUE_INTERVAL",
                            },
                        },
                    }
                }
            ) as _:
                pass

    def test_submit_run(self, instance, coordinator, workspace, remote_job):
        run = self.create_run_for_test(instance, remote_job, status=DagsterRunStatus.NOT_STARTED)
        returned_run = coordinator.submit_run(dg.SubmitRunContext(run, workspace))
        assert returned_run.run_id == run.run_id
        assert returned_run.status == DagsterRunStatus.QUEUED

        assert len(instance.run_launcher.queue()) == 0
        stored_run = instance.get_run_by_id(run.run_id)
        assert stored_run.status == DagsterRunStatus.QUEUED

        events = list(
            instance.get_records_for_run(
                run.run_id, of_type=DagsterEventType.PIPELINE_ENQUEUED
            ).records
        )
        assert len(events) == 1
        run_enqueued_data = events[0].event_log_entry.dagster_event.run_enqueued_data
        assert run_enqueued_data
        assert run_enqueued_data.code_location_name is not None
        assert run_enqueued_data.repository_name is not None

    def test_submit_run_checks_status(self, instance, coordinator, workspace, remote_job):
        run = self.create_run_for_test(instance, remote_job, status=DagsterRunStatus.QUEUED)
        coordinator.submit_run(dg.SubmitRunContext(run, workspace))

        # check that no enqueue event is reported (the submit run call is a no-op)
        assert (
            len(
                instance.get_records_for_run(
                    run.run_id, of_type=DagsterEventType.PIPELINE_ENQUEUED
                ).records
            )
            == 0
        )

    def test_cancel_run(self, instance, coordinator, workspace, remote_job):
        run = self.create_run_for_test(instance, remote_job, status=DagsterRunStatus.NOT_STARTED)

        coordinator.submit_run(dg.SubmitRunContext(run, workspace))

        coordinator.cancel_run(run.run_id)
        stored_run = instance.get_run_by_id(run.run_id)
        assert stored_run.status == DagsterRunStatus.CANCELED


def test_thread_config():
    num = 16
    with dg.instance_for_test(
        overrides={
            "run_coordinator": {
                "module": "dagster._core.run_coordinator",
                "class": "QueuedRunCoordinator",
                "config": {
                    "dequeue_use_threads": True,
                    "dequeue_num_workers": num,
                },
            }
        }
    ) as instance:
        assert instance.run_coordinator.dequeue_num_workers == num  # pyright: ignore[reportAttributeAccessIssue]
