# 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 abc
import json
from collections.abc import AsyncIterator, Hashable
from dataclasses import dataclass
from datetime import timedelta
from typing import TYPE_CHECKING, Annotated, Any

import structlog
from pydantic import (
    BaseModel,
    Discriminator,
    JsonValue,
    Tag,
    model_serializer,
)

from airflow.sdk.definitions._internal.templater import Templater
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.state import TaskInstanceState

log = structlog.get_logger(logger_name=__name__)

if TYPE_CHECKING:
    from typing import TypeAlias

    import jinja2

    from airflow.models.mappedoperator import MappedOperator
    from airflow.models.taskinstance import TaskInstance
    from airflow.sdk.definitions.context import Context
    from airflow.serialization.serialized_objects import SerializedBaseOperator

    Operator: TypeAlias = MappedOperator | SerializedBaseOperator


@dataclass
class StartTriggerArgs:
    """Arguments required for start task execution from triggerer."""

    trigger_cls: str
    next_method: str
    trigger_kwargs: dict[str, Any] | None = None
    next_kwargs: dict[str, Any] | None = None
    timeout: timedelta | None = None


class BaseTrigger(abc.ABC, Templater, LoggingMixin):
    """
    Base class for all triggers.

    A trigger has two contexts it can exist in:

     - Inside an Operator, when it's passed to TaskDeferred
     - Actively running in a trigger worker

    We use the same class for both situations, and rely on all Trigger classes
    to be able to return the arguments (possible to encode with Airflow-JSON) that will
    let them be re-instantiated elsewhere.
    """

    supports_triggerer_queue: bool = True

    def __init__(self, **kwargs):
        super().__init__()
        # these values are set by triggerer when preparing to run the instance
        # when run, they are injected into logger record.
        self._task_instance = None
        self.trigger_id = None
        self.template_fields = ()
        self.template_ext = ()
        self.task_id = None

    def _set_context(self, context):
        """Part of LoggingMixin and used mainly for configuration of task logging; not used for triggers."""
        pass

    @property
    def task(self) -> Operator | None:
        # We must check if the TaskInstance is the generated Pydantic one or the RuntimeTaskInstance
        if self.task_instance and hasattr(self.task_instance, "task"):
            return self.task_instance.task
        return None

    @property
    def task_instance(self) -> TaskInstance:
        return self._task_instance

    @task_instance.setter
    def task_instance(self, value: TaskInstance | None) -> None:
        self._task_instance = value
        if self.task_instance:
            self.task_id = self.task_instance.task_id
        if self.task:
            self.template_ext = self.task.template_ext
            # Only keep operator template_fields that are also keys in
            # start_trigger_args.trigger_kwargs *and* exist on the trigger.
            # Using the full operator template_fields would cause
            # AttributeError when the trigger does not have attributes with
            # the same names as the operator (e.g. "bash_command").
            #
            # When start_trigger_args is None (normal defer path), the triggerer
            # does not build a template context, so render_template_fields is
            # never called and empty template_fields is safe.
            start_trigger_args = getattr(self.task, "start_trigger_args", None)
            if start_trigger_args:
                from airflow.serialization.enums import Encoding

                raw = start_trigger_args.trigger_kwargs or {}
                # trigger_kwargs may be BaseSerialization-encoded; extract inner dict keys
                if isinstance(raw, dict) and Encoding.TYPE in raw:
                    raw = raw.get(Encoding.VAR) or {}
                trigger_kwarg_keys = set(raw.keys())
            else:
                trigger_kwarg_keys = set()
            if trigger_kwarg_keys:
                self.template_fields = tuple(
                    f for f in self.task.template_fields if f in trigger_kwarg_keys and hasattr(self, f)
                )
            else:
                self.template_fields = ()

    def render_template_fields(
        self,
        context: Context,
        jinja_env: jinja2.Environment | None = None,
    ) -> None:
        """
        Template all attributes listed in *self.template_fields*.

        This mutates the attributes in-place and is irreversible.

        :param context: Context dict with values to apply on content.
        :param jinja_env: Jinja's environment to use for rendering.
        """
        if not jinja_env:
            jinja_env = self.get_template_env()
        # self.template_fields is already filtered (in the task_instance setter) to only
        # include fields present in start_trigger_args.trigger_kwargs and on this trigger.
        self._do_render_template_fields(self, self.template_fields, context, jinja_env, set())

    @abc.abstractmethod
    def serialize(self) -> tuple[str, dict[str, Any]]:
        """
        Return the information needed to reconstruct this Trigger.

        :return: Tuple of (class path, keyword arguments needed to re-instantiate).
        """
        raise NotImplementedError("Triggers must implement serialize()")

    @abc.abstractmethod
    async def run(self) -> AsyncIterator[TriggerEvent]:
        """
        Run the trigger in an asynchronous context.

        The trigger should yield an Event whenever it wants to fire off
        an event, and return None if it is finished. Single-event triggers
        should thus yield and then immediately return.

        If it yields, it is likely that it will be resumed very quickly,
        but it may not be (e.g. if the workload is being moved to another
        triggerer process, or a multi-event trigger was being used for a
        single-event task defer).

        In either case, Trigger classes should assume they will be persisted,
        and then rely on cleanup() being called when they are no longer needed.
        """
        raise NotImplementedError("Triggers must implement run()")
        yield  # To convince Mypy this is an async iterator.

    async def cleanup(self) -> None:
        """
        Cleanup the trigger.

        Called when the trigger is no longer needed, and it's being removed
        from the active triggerer process.

        This method follows the async/await pattern to allow to run the cleanup
        in triggerer main event loop. Exceptions raised by the cleanup method
        are ignored, so if you would like to be able to debug them and be notified
        that cleanup method failed, you should wrap your code with try/except block
        and handle it appropriately (in async-compatible way).
        """

    async def on_kill(self) -> None:
        """
        Kill the external job managed by this trigger when the task is killed by a user.

        Symmetric with ``BaseOperator.on_kill()`` on the worker side: override this method
        to stop external work (e.g. cancel a BigQuery job, terminate a Databricks run) when
        a user explicitly acts on the deferred task via mark-failed, clear, or mark-succeeded.

        **Distinction from** ``cleanup()``:

        - ``cleanup()`` runs on every trigger exit — success, timeout, shutdown, and user
          kill. It is meant for releasing local resources held by this trigger instance.
          Putting external job cancellation in ``cleanup()`` would cancel in-flight work
          on every triggerer restart or rolling deploy.
        - ``on_kill()`` runs only when a user explicitly kills the task. It is the right
          place to cancel external work you do not want to keep running after the user performs an action.

        This only fires when a user acts on the task. It does not fire on:

        - Triggerer shutdown or restart — the trigger is redistributed, not cancelled.
        - Triggerer redistribution to another triggerer process.
        - Trigger timeout — the trigger is killed, not cancelled by user.
        - Normal trigger completion (the trigger fired an event).

        This method runs in the triggerer's asyncio event loop, so
        it must be async-safe. Use ``await`` for any I/O; do not block the event loop.

        Exceptions raised here are logged as warnings and do not
        propagate — they will not affect the task state or the triggerer. Implement your
        own retry or error handling inside this method if needed.

        ``on_kill()`` is given a bounded time to complete. Implementations
        that call slow external APIs should apply their own timeouts rather than relying on
        the framework bound.
        """

    @staticmethod
    def repr(classpath: str, kwargs: dict[str, Any]):
        kwargs_str = ", ".join(f"{k}={v}" for k, v in kwargs.items())
        return f"<{classpath} {kwargs_str}>"

    def __repr__(self) -> str:
        classpath, kwargs = self.serialize()
        return self.repr(classpath, kwargs)


class BaseEventTrigger(BaseTrigger):
    """
    Base class for triggers used to schedule DAGs based on external events.

    ``BaseEventTrigger`` is a subclass of ``BaseTrigger`` designed to identify triggers compatible with
    event-driven scheduling.

    **Sharing an underlying I/O stream between triggers**

    A subclass that polls an upstream resource which can be safely consumed
    by multiple sibling triggers (e.g. a directory scan, a polling REST API)
    may opt in to having the triggerer run a single underlying poll loop
    and fan its raw events out to every trigger in the group. To do so,
    override:

    * :meth:`shared_stream_key` — return a key identifying the
      shared stream (a ``tuple`` of strings is a common choice). Triggers
      whose key compares equal share one poll.
    * :meth:`open_shared_stream` — open the shared stream and yield raw
      events. Called once per group in the triggerer.
    * :meth:`filter_shared_stream` — convert the shared raw stream into this
      trigger's own ``TriggerEvent`` instances, applying any per-trigger
      filtering or transformation.

    Triggers whose ``shared_stream_key`` returns ``None`` (the default)
    keep the existing behavior: each trigger gets its own poll loop via
    :meth:`run`.

    **Suitable upstreams**

    The shared-stream channel is **one-way** today: events flow from the
    producer (``open_shared_stream``) to each subscriber's
    ``filter_shared_stream``, with no path back to tell the producer that a
    subscriber accepted, dropped, or finished processing an event. That
    restricts the pattern to upstreams whose consumption does **not** depend
    on a side effect on a handle that only the producer holds:

    * Idempotent / read-only reads (filesystem listings, polling REST APIs).
    * Subscriber-side-effect cleanup, where the trigger's per-event action
      (``unlink``, local marking, …) operates through APIs the subscriber
      already owns, independent of the shared producer handle.

    Upstreams **not** in scope include Kafka consumers (regardless of
    commit mode), SQS with delete-on-process or visibility extension,
    and any source where progress on the producer's handle is tied to
    the subscriber's accept / reject decision. These sources need a way
    for the subscriber to signal acceptance back to the producer, which
    the current shared-stream API does not provide.
    """

    supports_triggerer_queue: bool = False

    @staticmethod
    def hash(classpath: str, kwargs: dict[str, Any]) -> int:
        """
        Return the hash of the trigger classpath and kwargs. This is used to uniquely identify a trigger.

        We do not want to have this logic in ``BaseTrigger`` because, when used to defer tasks, 2 triggers
        can have the same classpath and kwargs. This is not true for event driven scheduling.
        """
        from airflow.serialization.encoders import encode_trigger
        from airflow.serialization.serialized_objects import BaseSerialization

        normalized = encode_trigger({"classpath": classpath, "kwargs": kwargs})["kwargs"]
        return hash((classpath, json.dumps(BaseSerialization.serialize(normalized)).encode("utf-8")))

    def shared_stream_key(self) -> Hashable | None:
        """
        Identify an underlying I/O stream that can be shared with sibling triggers.

        Two trigger instances whose ``shared_stream_key()`` return values
        compare equal (and are not ``None``) will share a single underlying
        poll loop in the triggerer. Each instance still receives the events
        it cares about through its own :meth:`filter_shared_stream` call.

        Returning ``None`` (the default) opts out of sharing — the trigger
        runs its own independent poll loop via :meth:`run`, exactly as today.

        The return value is read **once** when ``run_trigger`` first starts
        this trigger; any change to the key afterwards has no effect on
        group membership for this instance. To share one poll across a set
        of sibling triggers, ensure every trigger in the set returns the
        same key from the outset.

        The key must be deterministic — derive it from configuration fields,
        never from per-call values such as ``time.time()`` or ``uuid.uuid4()``,
        because the comparison must be stable across the lifetime of the group.

        .. note::

           This method is called **after** :meth:`render_template_fields`,
           so any templated attribute (for example a ``directory`` derived
           from a Jinja expression) is already resolved when the key is
           constructed. Two sibling triggers that render to the same path
           will correctly share their poll.
        """
        return None

    @classmethod
    async def open_shared_stream(cls, kwargs: dict[str, Any]) -> AsyncIterator[Any]:
        """
        Open the shared underlying stream and yield raw events.

        Called **once per shared-stream group** in the triggerer. ``kwargs``
        is taken from one trigger in the group; implementations should rely
        only on fields whose values participate in :meth:`shared_stream_key`,
        because other fields may differ between siblings in the group.

        Implementations are expected to run for the lifetime of the group —
        the triggerer drives the iterator from a single task and cancels it
        when the last subscriber leaves. Returning without raising (e.g.
        because the upstream resource closed) is treated as an error and
        propagated to every subscriber, so the contract is "yield forever, or
        raise". If an upstream EOF is a meaningful end-of-life condition,
        raise an exception that conveys it.

        Declared as a classmethod (not staticmethod) so subclasses can
        compose via ``super().open_shared_stream(kwargs)`` and reach
        ``cls`` for class-scoped state or diagnostics.

        Required only when :meth:`shared_stream_key` returns non-``None``.
        """
        raise NotImplementedError(
            f"{cls.__name__} declares a shared_stream_key but does not implement open_shared_stream"
        )
        yield  # pragma: no cover - convince mypy this is an async iterator

    async def filter_shared_stream(self, shared_stream: AsyncIterator[Any]) -> AsyncIterator[TriggerEvent]:
        """
        Transform the shared raw event stream into this trigger's events.

        The triggerer calls this method (instead of :meth:`run`) when this
        trigger participates in a shared-stream group. Iterate
        ``shared_stream`` to receive raw events from the shared poll, and
        ``yield`` a :class:`TriggerEvent` for each one that should fire this
        trigger.

        Required only when :meth:`shared_stream_key` returns non-``None``.
        """
        raise NotImplementedError(
            f"{type(self).__name__} declares a shared_stream_key but does not implement filter_shared_stream"
        )
        yield  # pragma: no cover - convince mypy this is an async iterator


class TriggerEvent(BaseModel):
    """
    Something that a trigger can fire when its conditions are met.

    Events must have a uniquely identifying value that would be the same
    wherever the trigger is run; this is to ensure that if the same trigger
    is being run in two locations (for HA reasons) that we can deduplicate its
    events.
    """

    payload: Any = None
    """
    The payload for the event to send back to the task.

    Must be natively JSON-serializable, or registered with the airflow serialization code.
    """

    def __init__(self, payload, **kwargs):
        super().__init__(payload=payload, **kwargs)

    def __repr__(self) -> str:
        return f"TriggerEvent<{self.payload!r}>"


class BaseTaskEndEvent(TriggerEvent):
    """
    Base event class to end the task without resuming on worker.

    :meta private:
    """

    task_instance_state: TaskInstanceState
    xcoms: dict[str, JsonValue] | None = None

    def __init__(self, *, xcoms: dict[str, JsonValue] | None = None, **kwargs) -> None:
        """
        Initialize the class with the specified parameters.

        :param xcoms: A dictionary of XComs or None.
        :param kwargs: Additional keyword arguments.
        """
        if "payload" in kwargs:
            raise ValueError("Param 'payload' not supported for this class.")
        # Yes this is _odd_. It's to support both constructor from users of
        # `TaskSuccessEvent(some_xcom_value)` and deserialization by pydantic.
        state = kwargs.pop("task_instance_state", self.__pydantic_fields__["task_instance_state"].default)
        super().__init__(payload=str(state), task_instance_state=state, **kwargs)
        self.xcoms = xcoms

    @model_serializer
    def ser_model(self) -> dict[str, Any]:
        # We need to customize the serialized schema so it works for the custom constructor we have to keep
        # the interface to this class "nice"
        return {"task_instance_state": self.task_instance_state, "xcoms": self.xcoms}


class TaskSuccessEvent(BaseTaskEndEvent):
    """Yield this event in order to end the task successfully."""

    task_instance_state: TaskInstanceState = TaskInstanceState.SUCCESS


class TaskFailedEvent(BaseTaskEndEvent):
    """Yield this event in order to end the task with failure."""

    task_instance_state: TaskInstanceState = TaskInstanceState.FAILED


class TaskSkippedEvent(BaseTaskEndEvent):
    """Yield this event in order to end the task with status 'skipped'."""

    task_instance_state: TaskInstanceState = TaskInstanceState.SKIPPED


def trigger_event_discriminator(v):
    if isinstance(v, dict):
        return v.get("task_instance_state", "_event_")
    if isinstance(v, TriggerEvent):
        return getattr(v, "task_instance_state", "_event_")


DiscrimatedTriggerEvent = Annotated[
    Annotated[TriggerEvent, Tag("_event_")]
    | Annotated[TaskSuccessEvent, Tag(TaskInstanceState.SUCCESS)]
    | Annotated[TaskFailedEvent, Tag(TaskInstanceState.FAILED)]
    | Annotated[TaskSkippedEvent, Tag(TaskInstanceState.SKIPPED)],
    Discriminator(trigger_event_discriminator),
]
