#
# 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 datetime
import logging
from abc import abstractmethod
from collections.abc import (
    Callable,
    Collection,
    Iterable,
    Iterator,
)
from typing import TYPE_CHECKING, Any, TypeAlias

from airflow.sdk import TriggerRule, WeightRule
from airflow.sdk.configuration import conf
from airflow.sdk.definitions._internal.mixins import DependencyMixin
from airflow.sdk.definitions._internal.node import DAGNode
from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext
from airflow.sdk.definitions._internal.templater import Templater
from airflow.sdk.definitions.context import Context

if TYPE_CHECKING:
    import jinja2

    from airflow.sdk.bases.operator import BaseOperator
    from airflow.sdk.bases.operatorlink import BaseOperatorLink
    from airflow.sdk.definitions.dag import DAG
    from airflow.sdk.definitions.mappedoperator import MappedOperator
    from airflow.sdk.definitions.taskgroup import MappedTaskGroup

TaskStateChangeCallback = Callable[[Context], None]
TaskStateChangeCallbackAttrType: TypeAlias = TaskStateChangeCallback | list[TaskStateChangeCallback] | None

DEFAULT_OWNER: str = conf.get_mandatory_value("operators", "default_owner")
DEFAULT_POOL_SLOTS: int = 1
DEFAULT_POOL_NAME = "default_pool"
DEFAULT_PRIORITY_WEIGHT: int = 1
# Databases do not support arbitrary precision integers, so we need to limit the range of priority weights.
# postgres: -2147483648 to +2147483647 (see https://www.postgresql.org/docs/current/datatype-numeric.html)
# mysql: -2147483648 to +2147483647 (see https://dev.mysql.com/doc/refman/8.4/en/integer-types.html)
# sqlite: -9223372036854775808 to +9223372036854775807 (see https://sqlite.org/datatype3.html)
MINIMUM_PRIORITY_WEIGHT: int = -2147483648
MAXIMUM_PRIORITY_WEIGHT: int = 2147483647
DEFAULT_EXECUTOR: str | None = None
DEFAULT_QUEUE: str = conf.get_mandatory_value("operators", "default_queue")
DEFAULT_IGNORE_FIRST_DEPENDS_ON_PAST: bool = False
DEFAULT_WAIT_FOR_PAST_DEPENDS_BEFORE_SKIPPING: bool = False
DEFAULT_RETRIES: int = conf.getint("core", "default_task_retries", fallback=0)
DEFAULT_RETRY_DELAY: datetime.timedelta = datetime.timedelta(
    seconds=conf.getint("core", "default_task_retry_delay", fallback=300)
)
MAX_RETRY_DELAY: int = conf.getint("core", "max_task_retry_delay", fallback=24 * 60 * 60)

# TODO: Task-SDK -- these defaults should be overridable from the Airflow config
DEFAULT_TRIGGER_RULE: TriggerRule = TriggerRule.ALL_SUCCESS
DEFAULT_WEIGHT_RULE: WeightRule = WeightRule(
    conf.get("core", "default_task_weight_rule", fallback=WeightRule.DOWNSTREAM)
)
DEFAULT_TASK_EXECUTION_TIMEOUT: datetime.timedelta | None = conf.gettimedelta(
    "core", "default_task_execution_timeout"
)
DEFAULT_EMAIL_ON_FAILURE: bool = conf.getboolean("email", "default_email_on_failure", fallback=True)
DEFAULT_EMAIL_ON_RETRY: bool = conf.getboolean("email", "default_email_on_retry", fallback=True)
log = logging.getLogger(__name__)


class AbstractOperator(Templater, DAGNode):
    """
    Common implementation for operators, including unmapped and mapped.

    This base class is more about sharing implementations, not defining a common
    interface. Unfortunately it's difficult to use this as the common base class
    for typing due to BaseOperator carrying too much historical baggage.

    The union type ``from airflow.models.operator import Operator`` is easier
    to use for typing purposes.

    :meta private:
    """

    operator_class: type[BaseOperator] | dict[str, Any]

    priority_weight: int

    # Defines the operator level extra links.
    operator_extra_links: Collection[BaseOperatorLink] = ()

    owner: str
    task_id: str

    outlets: list
    inlets: list

    trigger_rule: TriggerRule
    _needs_expansion: bool | None = None
    _on_failure_fail_dagrun = False
    is_setup: bool = False
    is_teardown: bool = False
    returns_dag_result: bool = False

    @property
    def is_async(self) -> bool:
        return False

    @property
    def task_type(self) -> str:
        raise NotImplementedError()

    @property
    def operator_name(self) -> str:
        raise NotImplementedError()

    _is_sensor: bool = False
    _is_mapped: bool = False
    _can_skip_downstream: bool = False

    @property
    def dag_id(self) -> str:
        """Returns dag id if it has one or an adhoc + owner."""
        dag = self.get_dag()
        if dag:
            return dag.dag_id
        return f"adhoc_{self.owner}"

    @property
    def node_id(self) -> str:
        return self.task_id

    @property
    @abstractmethod
    def task_display_name(self) -> str: ...

    @property
    def is_mapped(self):
        return self._is_mapped

    @property
    def label(self) -> str | None:
        if self.task_display_name and self.task_display_name != self.task_id:
            return self.task_display_name
        # Prefix handling if no display is given is cloned from taskmixin for compatibility
        tg = self.task_group
        if tg and tg.node_id and tg.prefix_group_id:
            # "task_group_id.task_id" -> "task_id"
            return self.task_id[len(tg.node_id) + 1 :]
        return self.task_id

    @property
    def on_failure_fail_dagrun(self):
        """
        Whether the operator should fail the dagrun on failure.

        :meta private:
        """
        return self._on_failure_fail_dagrun

    @on_failure_fail_dagrun.setter
    def on_failure_fail_dagrun(self, value):
        """
        Setter for on_failure_fail_dagrun property.

        :meta private:
        """
        if value is True and self.is_teardown is not True:
            raise ValueError(
                f"Cannot set task on_failure_fail_dagrun for "
                f"'{self.task_id}' because it is not a teardown task."
            )
        self._on_failure_fail_dagrun = value

    @property
    def inherits_from_empty_operator(self):
        """Used to determine if an Operator is inherited from EmptyOperator."""
        # This looks like `isinstance(self, EmptyOperator) would work, but this also
        # needs to cope when `self` is a Serialized instance of a EmptyOperator or one
        # of its subclasses (which don't inherit from anything but BaseOperator).
        return getattr(self, "_is_empty", False)

    @property
    def inherits_from_skipmixin(self):
        """Used to determine if an Operator is inherited from SkipMixin or its subclasses (e.g., BranchMixin)."""
        return getattr(self, "_can_skip_downstream", False)

    def as_setup(self):
        self.is_setup = True
        return self

    def as_teardown(
        self,
        *,
        setups: BaseOperator | Iterable[BaseOperator] | None = None,
        on_failure_fail_dagrun: bool | None = None,
    ):
        self.is_teardown = True
        self.trigger_rule = TriggerRule.ALL_DONE_SETUP_SUCCESS
        if on_failure_fail_dagrun is not None:
            self.on_failure_fail_dagrun = on_failure_fail_dagrun
        if setups is not None:
            setups = [setups] if isinstance(setups, DependencyMixin) else setups
            for s in setups:
                s.is_setup = True
                s >> self
        return self

    def __enter__(self):
        if not self.is_setup and not self.is_teardown:
            raise RuntimeError("Only setup/teardown tasks can be used as context managers.")
        SetupTeardownContext.push_setup_teardown_task(self)
        return SetupTeardownContext

    def __exit__(self, exc_type, exc_val, exc_tb):
        SetupTeardownContext.set_work_task_roots_and_leaves()

    # TODO: Task-SDK -- Should the following methods removed?
    #   get_template_env
    #   _render
    def get_template_env(self, dag: DAG | None = None) -> jinja2.Environment:
        """Get the template environment for rendering templates."""
        from airflow.sdk.definitions._internal.templater import create_template_env

        if dag is None:
            dag = self.get_dag()
        # Check if the operator has an explicit native rendering preference
        render_op_template_as_native_obj = getattr(self, "render_template_as_native_obj", None)
        if render_op_template_as_native_obj is not None:
            if dag:
                # Use dag's template settings (searchpath, macros, filters, etc.)
                searchpath = [dag.folder]
                if dag.template_searchpath:
                    searchpath += dag.template_searchpath
                return create_template_env(
                    native=render_op_template_as_native_obj,
                    searchpath=searchpath,
                    template_undefined=dag.template_undefined,
                    jinja_environment_kwargs=dag.jinja_environment_kwargs,
                    user_defined_macros=dag.user_defined_macros,
                    user_defined_filters=dag.user_defined_filters,
                )
            # No dag context available, use minimal template env
            return create_template_env(native=render_op_template_as_native_obj)
        # No operator-level override, delegate to parent class
        return super().get_template_env(dag=dag)

    def _render(self, template, context, dag: DAG | None = None):
        if dag is None:
            dag = self.get_dag()
        return super()._render(template, context, dag=dag)

    def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]:
        """
        Return mapped nodes that are direct dependencies of the current task.

        For now, this walks the entire Dag to find mapped nodes that has this
        current task as an upstream. We cannot use ``downstream_list`` since it
        only contains operators, not task groups. In the future, we should
        provide a way to record a Dag node's all downstream nodes instead.

        Note that this does not guarantee the returned tasks actually use the
        current task for task mapping, but only checks those task are mapped
        operators, and are downstreams of the current task.

        To get a list of tasks that uses the current task for task mapping, use
        :meth:`iter_mapped_dependants` instead.
        """
        from airflow.sdk.definitions.mappedoperator import MappedOperator
        from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup

        def _walk_group(group: TaskGroup) -> Iterable[tuple[str, DAGNode]]:
            """
            Recursively walk children in a task group.

            This yields all direct children (including both tasks and task
            groups), and all children of any task groups.
            """
            for key, child in group.children.items():
                yield key, child
                if isinstance(child, TaskGroup):
                    yield from _walk_group(child)

        dag = self.get_dag()
        if not dag:
            raise RuntimeError("Cannot check for mapped dependants when not attached to a Dag")
        for key, child in _walk_group(dag.task_group):
            if key == self.node_id:
                continue
            if not isinstance(child, (MappedOperator, MappedTaskGroup)):
                continue
            if self.node_id in child.upstream_task_ids:
                yield child

    def iter_mapped_dependants(self) -> Iterator[MappedOperator | MappedTaskGroup]:
        """
        Return mapped nodes that depend on the current task the expansion.

        For now, this walks the entire Dag to find mapped nodes that has this
        current task as an upstream. We cannot use ``downstream_list`` since it
        only contains operators, not task groups. In the future, we should
        provide a way to record a Dag node's all downstream nodes instead.
        """
        return (
            downstream
            for downstream in self._iter_all_mapped_downstreams()
            if any(p.node_id == self.node_id for p in downstream.iter_mapped_dependencies())
        )

    def iter_mapped_task_groups(self) -> Iterator[MappedTaskGroup]:
        """
        Return mapped task groups this task belongs to.

        Groups are returned from the innermost to the outmost.

        :meta private:
        """
        if (group := self.task_group) is None:
            return
        yield from group.iter_mapped_task_groups()

    def get_closest_mapped_task_group(self) -> MappedTaskGroup | None:
        """
        Get the mapped task group "closest" to this task in the Dag.

        :meta private:
        """
        return next(self.iter_mapped_task_groups(), None)

    def get_needs_expansion(self) -> bool:
        """
        Return true if the task is MappedOperator or is in a mapped task group.

        :meta private:
        """
        if self._needs_expansion is None:
            if self.get_closest_mapped_task_group() is not None:
                self._needs_expansion = True
            else:
                self._needs_expansion = False
        return self._needs_expansion
