# 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 asyncio
import logging
from collections.abc import Sequence
from functools import partial
from typing import Any

from asgiref.sync import sync_to_async

from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_0_PLUS
from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import AirflowException
from airflow.triggers.base import TriggerEvent

if AIRFLOW_V_3_0_PLUS:
    from airflow.triggers.base import BaseEventTrigger
else:
    from airflow.triggers.base import BaseTrigger as BaseEventTrigger  # type: ignore

log = logging.getLogger(__name__)


class AwaitMessageTrigger(BaseEventTrigger):
    """
    A trigger that waits for a message matching specific criteria to arrive in Kafka.

    The behavior of the consumer of this trigger is as follows:
    - poll the Kafka topics for a message, if no message returned, sleep
    - process the message with provided callable:

        - if callable is provided and returns any data, raise a TriggerEvent with the return data
        - else raise a TriggerEvent with the original message

    - by default, the message offset is committed after processing. This can be
      disabled by setting ``commit_offset=False``, allowing manual offset management
      in downstream tasks.

    :param kafka_config_id: The connection object to use, defaults to "kafka_default"
    :param topics: The topic (or topic regex) that should be searched for messages
    :param apply_function: the location of the function to apply to messages for determination of matching
        criteria. (In python dot notation as a string)
    :param apply_function_args: A set of arguments to apply to the callable, defaults to None
    :param apply_function_kwargs: A set of key word arguments to apply to the callable, defaults to None,
        defaults to None
    :param poll_timeout: How long the Kafka client should wait before returning from a poll request to
        Kafka (seconds), defaults to 1
    :param poll_interval: How long the trigger should sleep after reaching the end of the Kafka log
        (seconds), defaults to 5
    :param commit_offset: Whether to commit the message offset after poll.
        If set to False, the offset is not committed automatically, allowing
        downstream tasks to handle offset committing manually (e.g., after
        successful processing). Defaults to True.

    """

    def __init__(
        self,
        topics: Sequence[str],
        apply_function: str | None = None,
        kafka_config_id: str = "kafka_default",
        apply_function_args: Sequence[Any] | None = None,
        apply_function_kwargs: dict[Any, Any] | None = None,
        poll_timeout: float = 1,
        poll_interval: float = 5,
        commit_offset: bool = True,
    ) -> None:
        super().__init__()
        self.topics = topics
        self.apply_function = apply_function
        self.apply_function_args = apply_function_args or ()
        self.apply_function_kwargs = apply_function_kwargs or {}
        self.kafka_config_id = kafka_config_id
        self.poll_timeout = poll_timeout
        self.poll_interval = poll_interval
        self.commit_offset = commit_offset
        self._consumer = None

    def serialize(self) -> tuple[str, dict[str, Any]]:
        return (
            "airflow.providers.apache.kafka.triggers.await_message.AwaitMessageTrigger",
            {
                "topics": self.topics,
                "apply_function": self.apply_function,
                "apply_function_args": self.apply_function_args,
                "apply_function_kwargs": self.apply_function_kwargs,
                "kafka_config_id": self.kafka_config_id,
                "poll_timeout": self.poll_timeout,
                "poll_interval": self.poll_interval,
                "commit_offset": self.commit_offset,
            },
        )

    async def run(self):
        consumer_hook = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id)

        async_get_consumer = sync_to_async(consumer_hook.get_consumer)
        self._consumer = await async_get_consumer()

        async_poll = sync_to_async(self._consumer.poll)
        async_commit = sync_to_async(self._consumer.commit)

        async_message_process = None
        if self.apply_function:
            processing_call = import_string(self.apply_function)
            processing_call = partial(
                processing_call, *self.apply_function_args, **self.apply_function_kwargs
            )
            async_message_process = sync_to_async(processing_call)
        while True:
            message = await async_poll(self.poll_timeout)

            if message is None:
                continue
            elif message.error():
                raise AirflowException(f"Error: {message.error()}")
            else:
                event = (
                    await async_message_process(message)
                    if async_message_process
                    else message.value().decode("utf-8")
                )
                if event:
                    if self.commit_offset:
                        await async_commit(message=message, asynchronous=False)
                    yield TriggerEvent(event)
                    break
                else:
                    if self.commit_offset:
                        await async_commit(message=message, asynchronous=False)
                    await asyncio.sleep(self.poll_interval)

    async def cleanup(self) -> None:
        consumer = self._consumer
        if consumer is not None:
            self._consumer = None
            try:
                await sync_to_async(consumer.close)()
            except Exception:
                log.warning("Failed to close Kafka consumer", exc_info=True)
