import logging
from collections.abc import Mapping, Sequence
from typing import Any, NamedTuple, Optional

from typing_extensions import Self

from dagster import (
    DagsterEvent,
    IntSource,
    String,
    _check as check,
)
from dagster._builtins import Bool
from dagster._config import Array, Field, Noneable, ScalarUnion, Shape
from dagster._config.config_schema import UserConfigSchema
from dagster._core.instance import T_DagsterInstance
from dagster._core.run_coordinator.base import RunCoordinator, SubmitRunContext
from dagster._core.storage.dagster_run import DagsterRun, DagsterRunStatus
from dagster._serdes import ConfigurableClass, ConfigurableClassData


class RunQueueConfig(
    NamedTuple(
        "_RunQueueConfig",
        [
            ("max_concurrent_runs", int),
            ("tag_concurrency_limits", Sequence[Mapping[str, Any]]),
            ("max_user_code_failure_retries", int),
            ("user_code_failure_retry_delay", int),
            ("should_block_op_concurrency_limited_runs", bool),
            ("op_concurrency_slot_buffer", int),
            ("explicitly_enabled_concurrency_run_blocking", bool),
        ],
    )
):
    def __new__(
        cls,
        max_concurrent_runs: int,
        tag_concurrency_limits: Optional[Sequence[Mapping[str, Any]]],
        max_user_code_failure_retries: int = 0,
        user_code_failure_retry_delay: int = 60,
        should_block_op_concurrency_limited_runs: bool = True,
        op_concurrency_slot_buffer: int = 0,
        explicitly_enabled_concurrency_run_blocking: bool = False,
    ):
        return super().__new__(
            cls,
            check.int_param(max_concurrent_runs, "max_concurrent_runs"),
            check.opt_sequence_param(tag_concurrency_limits, "tag_concurrency_limits"),
            check.int_param(max_user_code_failure_retries, "max_user_code_failure_retries"),
            check.int_param(user_code_failure_retry_delay, "user_code_failure_retry_delay"),
            check.bool_param(
                should_block_op_concurrency_limited_runs, "should_block_op_concurrency_limited_runs"
            ),
            check.int_param(op_concurrency_slot_buffer, "op_concurrency_slot_buffer"),
            check.bool_param(
                explicitly_enabled_concurrency_run_blocking,
                "explicitly_enabled_concurrency_run_blocking",
            ),
        )

    def with_concurrency_settings(
        self, concurrency_settings: Mapping[str, Any]
    ) -> "RunQueueConfig":
        run_settings = concurrency_settings.get("runs", {})
        pool_settings = concurrency_settings.get("pools", {})
        return RunQueueConfig(
            max_concurrent_runs=run_settings.get("max_concurrent_runs", self.max_concurrent_runs),
            tag_concurrency_limits=run_settings.get(
                "tag_concurrency_limits", self.tag_concurrency_limits
            ),
            max_user_code_failure_retries=self.max_user_code_failure_retries,
            user_code_failure_retry_delay=self.user_code_failure_retry_delay,
            should_block_op_concurrency_limited_runs=self.should_block_op_concurrency_limited_runs,
            op_concurrency_slot_buffer=pool_settings.get(
                "op_granularity_run_buffer", self.op_concurrency_slot_buffer
            ),
            explicitly_enabled_concurrency_run_blocking=self.explicitly_enabled_concurrency_run_blocking,
        )


class QueuedRunCoordinator(RunCoordinator[T_DagsterInstance], ConfigurableClass):
    """Enqueues runs via the run storage, to be deqeueued by the Dagster Daemon process. Requires
    the Dagster Daemon process to be alive in order for runs to be launched.
    """

    def __init__(
        self,
        max_concurrent_runs: Optional[int] = None,
        tag_concurrency_limits: Optional[Sequence[Mapping[str, Any]]] = None,
        dequeue_interval_seconds: Optional[int] = None,
        dequeue_use_threads: Optional[bool] = None,
        dequeue_num_workers: Optional[int] = None,
        max_user_code_failure_retries: Optional[int] = None,
        user_code_failure_retry_delay: Optional[int] = None,
        block_op_concurrency_limited_runs: Optional[Mapping[str, Any]] = None,
        inst_data: Optional[ConfigurableClassData] = None,
    ):
        self._inst_data: Optional[ConfigurableClassData] = check.opt_inst_param(
            inst_data, "inst_data", ConfigurableClassData
        )
        self._max_concurrent_runs: int = check.opt_int_param(
            max_concurrent_runs, "max_concurrent_runs", 10
        )
        check.invariant(
            self._max_concurrent_runs >= -1,
            "Negative values other than -1 (which disables the limit) for max_concurrent_runs"
            " are disallowed.",
        )
        self._tag_concurrency_limits: Sequence[Mapping[str, Any]] = check.opt_list_param(
            tag_concurrency_limits,
            "tag_concurrency_limits",
        )
        self._dequeue_interval_seconds: int = check.opt_int_param(
            dequeue_interval_seconds, "dequeue_interval_seconds", 5
        )
        self._dequeue_use_threads: bool = check.opt_bool_param(
            dequeue_use_threads, "dequeue_use_threads", False
        )
        self._dequeue_num_workers: Optional[int] = check.opt_int_param(
            dequeue_num_workers, "dequeue_num_workers"
        )
        self._max_user_code_failure_retries: int = check.opt_int_param(
            max_user_code_failure_retries, "max_user_code_failure_retries", 0
        )
        self._user_code_failure_retry_delay: int = check.opt_int_param(
            user_code_failure_retry_delay, "user_code_failure_retry_delay", 60
        )
        self._should_block_op_concurrency_limited_runs: bool = bool(
            not block_op_concurrency_limited_runs
            or block_op_concurrency_limited_runs.get("enabled", True)
        )
        self._explicitly_enabled_concurrency_run_blocking: bool = bool(
            block_op_concurrency_limited_runs and block_op_concurrency_limited_runs.get("enabled")
        )
        self._op_concurrency_slot_buffer: int = (
            block_op_concurrency_limited_runs.get("op_concurrency_slot_buffer", 0)
            if block_op_concurrency_limited_runs
            else 0
        )
        if self._op_concurrency_slot_buffer:
            check.invariant(
                self._should_block_op_concurrency_limited_runs,
                "op_concurrency_slot_buffer can only be set if block_op_concurrency_limited_runs "
                "is enabled",
            )
        self._logger = logging.getLogger("dagster.run_coordinator.queued_run_coordinator")
        super().__init__()

    @property
    def inst_data(self) -> Optional[ConfigurableClassData]:
        return self._inst_data

    def get_run_queue_config(self) -> RunQueueConfig:
        return RunQueueConfig(
            max_concurrent_runs=self._max_concurrent_runs,
            tag_concurrency_limits=self._tag_concurrency_limits,
            max_user_code_failure_retries=self._max_user_code_failure_retries,
            user_code_failure_retry_delay=self._user_code_failure_retry_delay,
            should_block_op_concurrency_limited_runs=self._should_block_op_concurrency_limited_runs,
            op_concurrency_slot_buffer=self._op_concurrency_slot_buffer,
            explicitly_enabled_concurrency_run_blocking=self._explicitly_enabled_concurrency_run_blocking,
        )

    @property
    def dequeue_interval_seconds(self) -> int:
        return self._dequeue_interval_seconds

    @property
    def dequeue_use_threads(self) -> bool:
        return self._dequeue_use_threads

    @property
    def dequeue_num_workers(self) -> Optional[int]:
        return self._dequeue_num_workers

    @property
    def should_block_op_concurrency_limited_runs(self) -> bool:
        return self._should_block_op_concurrency_limited_runs

    @property
    def op_concurrency_slot_buffer(self) -> int:
        return self._op_concurrency_slot_buffer

    @classmethod
    def config_type(cls) -> UserConfigSchema:
        return {
            "max_concurrent_runs": Field(
                config=IntSource,
                is_required=False,
                description=(
                    "The maximum number of runs that are allowed to be in progress at once."
                    " Defaults to 10. Set to -1 to disable the limit. Set to 0 to stop any runs"
                    " from launching. Any other negative values are disallowed."
                ),
            ),
            "tag_concurrency_limits": Field(
                config=Noneable(
                    Array(
                        Shape(
                            {
                                "key": String,
                                "value": Field(
                                    ScalarUnion(
                                        scalar_type=String,
                                        non_scalar_schema=Shape({"applyLimitPerUniqueValue": Bool}),
                                    ),
                                    is_required=False,
                                ),
                                "limit": Field(int),
                            }
                        )
                    )
                ),
                is_required=False,
                description=(
                    "A set of limits that are applied to runs with particular tags. If a value is"
                    " set, the limit is applied to only that key-value pair. If no value is set,"
                    " the limit is applied across all values of that key. If the value is set to a"
                    " dict with `applyLimitPerUniqueValue: true`, the limit will apply to the"
                    " number of unique values for that key."
                ),
            ),
            "dequeue_interval_seconds": Field(
                config=IntSource,
                is_required=False,
                description=(
                    "The interval in seconds at which the Dagster Daemon "
                    "should periodically check the run queue for new runs to launch."
                ),
            ),
            "dequeue_use_threads": Field(
                config=bool,
                is_required=False,
                description=(
                    "Whether or not to use threads for concurrency when launching dequeued runs."
                ),
            ),
            "dequeue_num_workers": Field(
                config=IntSource,
                is_required=False,
                description=(
                    "If dequeue_use_threads is true, limit the number of concurrent worker threads."
                ),
            ),
            "max_user_code_failure_retries": Field(
                config=IntSource,
                is_required=False,
                default_value=0,
                description=(
                    "If there is an error reaching a Dagster gRPC server while dequeuing the run,"
                    " how many times to retry the dequeue before failing it. The only run launcher"
                    " that requires the gRPC server to be running is the DefaultRunLauncher, so"
                    " setting this will have no effect unless that run launcher is being used."
                ),
            ),
            "user_code_failure_retry_delay": Field(
                config=IntSource,
                is_required=False,
                default_value=60,
                description=(
                    "If there is an error reaching a Dagster gRPC server while dequeuing the run,"
                    " how long to wait before retrying any runs from that same code location. The"
                    " only run launcher that requires the gRPC server to be running is the"
                    " DefaultRunLauncher, so setting this will have no effect unless that run"
                    " launcher is being used."
                ),
            ),
            "block_op_concurrency_limited_runs": Field(
                {
                    "enabled": Field(Bool, is_required=False),
                    "op_concurrency_slot_buffer": Field(
                        int,
                        is_required=False,
                        description=(
                            "Determines whether or not a run will be dequeued if it consists of ops that "
                            "will all be initially blocked waiting for global op concurrency slots to be "
                            "free."
                        ),
                    ),
                }
            ),
        }

    @classmethod
    def from_config_value(
        cls, inst_data: ConfigurableClassData, config_value: Mapping[str, Any]
    ) -> Self:
        return cls(
            inst_data=inst_data,
            max_concurrent_runs=config_value.get("max_concurrent_runs"),
            tag_concurrency_limits=config_value.get("tag_concurrency_limits"),
            dequeue_interval_seconds=config_value.get("dequeue_interval_seconds"),
            dequeue_use_threads=config_value.get("dequeue_use_threads"),
            dequeue_num_workers=config_value.get("dequeue_num_workers"),
            max_user_code_failure_retries=config_value.get("max_user_code_failure_retries"),
            user_code_failure_retry_delay=config_value.get("user_code_failure_retry_delay"),
            block_op_concurrency_limited_runs=config_value.get("block_op_concurrency_limited_runs"),
        )

    def submit_run(self, context: SubmitRunContext) -> DagsterRun:
        dagster_run = context.dagster_run

        if dagster_run.status == DagsterRunStatus.NOT_STARTED:
            enqueued_event = DagsterEvent.job_enqueue(dagster_run)
            self._instance.report_dagster_event(enqueued_event, run_id=dagster_run.run_id)
        else:
            # the run was already submitted, this is a no-op
            self._logger.warning(
                f"submit_run called for run {dagster_run.run_id} with status "
                f"{dagster_run.status.value}, skipping enqueue."
            )

        run = self._instance.get_run_by_id(dagster_run.run_id)
        if run is None:
            check.failed(f"Failed to reload run {dagster_run.run_id}")
        assert run
        return run

    def cancel_run(self, run_id: str) -> bool:
        run = self._instance.get_run_by_id(run_id)
        if not run:
            return False
        # NOTE: possible race condition if the dequeuer acts on this run at the same time
        # https://github.com/dagster-io/dagster/issues/3323
        if run.status == DagsterRunStatus.QUEUED:
            self._instance.report_run_canceling(
                run,
                message="Canceling run from the queue.",
            )
            self._instance.report_run_canceled(run)
            return True
        else:
            return self._instance.run_launcher.terminate(run_id)
