#  Copyright (c) ZenML GmbH 2026. All Rights Reserved.
#
#  Licensed 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:
#
#       https://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.
"""Base abstractions for the pub/sub layer."""

import asyncio
import logging
import random
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Any

from pydantic import Field

from zenml.constants import (
    ENV_ZENML_CONSUMER_POLLING_PREFIX,
    ENV_ZENML_CONSUMER_PREFIX,
    ENV_ZENML_PRODUCER_PREFIX,
)
from zenml.pub_sub.models import (
    CriticalEvent,
    CriticalEventType,
    MessageEnvelope,
    MessagePayload,
)
from zenml.utils.env_utils import ConfigBase

logger = logging.getLogger(__name__)


class MessageDecodeError(Exception):
    """Raised when a raw broker message cannot be decoded/preprocessed."""


class MessageExecutionError(Exception):
    """Raised when message processing fails after retries."""


class MessageEncodeError(Exception):
    """Raised when a payload cannot be encoded into a broker message."""


class MessageSubmissionError(Exception):
    """Raised when sending a message fails."""


class ProducerConfig(ConfigBase):
    """Config object grouping producer configuration params."""

    @staticmethod
    def prefixes() -> list[str]:
        """Implementation of the prefixes config method.

        Returns:
            A list of prefixes to use for env var name resolution.
        """
        return [ENV_ZENML_PRODUCER_PREFIX]

    send_retries: int = Field(
        default=1,
        description="Number of retries to send messages before failing.",
        ge=0,
    )


class CriticalEventHandler(ABC):
    """Abstraction for handling critical events."""

    @abstractmethod
    async def handle(self, event: CriticalEvent) -> None:
        """Handles a critical event.

        Handling of a critical event would be include the following operations:
        - Tracking of the event instance and their evolution
        - Alerting mechanisms based on event monitoring alarms
        - Mitigation actions

        Args:
            event: A CriticalEvent instance.
        """
        pass


class ProducerBase(ABC):
    """Async producer base.

    Backends implement:
      1) build_message(payload) -> backend-specific message object (Any)
      2) send_message(message) -> submit to broker

    Base provides:
      - publish(payload): build + send (+ optional retries) and emits CriticalEvent on failure.
    """

    def __init__(
        self,
        config: ProducerConfig,
        event_handler: CriticalEventHandler | None = None,
    ) -> None:
        """Producer constructor.

        Args:
            config: The producer config object.
            event_handler: A CriticalEventHandler instance (optional).
        """
        self._cfg = config
        self._event_handler = event_handler

    @property
    def config(self) -> ProducerConfig:
        """Implements the 'config' property.

        Returns:
            The producer config object.
        """
        return self._cfg

    @abstractmethod
    async def build_message(self, payload: MessagePayload) -> Any:
        """Convert the message to a broker-specific payload.

        Args:
            payload: A MessagePayload object.

        Returns:
            A broker-specific message.
        """

    @abstractmethod
    async def send_message(self, message: Any) -> str:
        """Sends a message to the broker.

        Args:
            message: A broker-specific message.

        Returns:
            The message ID generated by the broker.
        """

    async def handle_critical_event(self, event: CriticalEvent) -> None:
        """Critical event handler.

        A critical event indicates that the application behavior has degraded.
        For instance, the inability to serialize the message to a broker-compatible format
        is extremely severe as it most probably repeats across messages meaning no messages
        are published for execution at all.

        A complete implementation of a critical event handling process would be the following:

        1. Log in detailed fashion the event so that it is traceable in the logs.
        2. Publish the event to an event aggregator service so that it is tracked long-term and
        alerts can be generated based on its frequency and severity.
        3. Hold the message in a DLQ to enable re-execution when a bug fix is introduced in the code.

        Args:
            event: A critical event instance.
        """
        logger.exception(
            "Critical event %s caused by the following exception:",
            event.value,
            exc_info=event.exception if event.exception else None,
        )
        if self._event_handler:
            asyncio.create_task(self._event_handler.handle(event))

    async def publish(self, payload: MessagePayload) -> str:
        """Build and send message, emitting CriticalEvent on failure.

        Args:
            payload: Application payload.

        Returns:
            Message ID generated by the broker.

        Raises:
            MessageEncodeError: If encoding/building fails.
            MessageSubmissionError: If sending fails after retries.
        """
        try:
            message = await self.build_message(payload)
        except Exception as exc:
            await self.handle_critical_event(
                CriticalEvent(
                    value=CriticalEventType.ENCODE_FAILED,
                    description=f"Encode failed for payload {payload.id}: {exc!r}",
                    created_at=datetime.now(tz=timezone.utc),
                    exception=exc,
                )
            )
            raise MessageEncodeError(
                f"Encode failed for payload {payload.id}"
            ) from exc

        attempts = self.config.send_retries + 1
        last_exc = None

        for attempt in range(attempts):
            try:
                return await self.send_message(message)
            except Exception as exc:
                last_exc = exc
                logger.exception(
                    "Attempt %s to send message %s failed",
                    attempt,
                    payload.id,
                )

        await self.handle_critical_event(
            CriticalEvent(
                value=CriticalEventType.SEND_FAILED,
                description=f"Submission failed for payload {payload.id}",
                created_at=datetime.now(tz=timezone.utc),
                exception=last_exc,
            )
        )
        raise MessageSubmissionError(f"Send failed for payload {payload.id}")


class ConsumerRuntimeConfig(ConfigBase):
    """Config shared by polling and push consumers."""

    @staticmethod
    def prefixes() -> list[str]:
        """Implementation of the prefixes config method.

        Returns:
            A list of prefixes to use for env var name resolution.
        """
        return [ENV_ZENML_CONSUMER_PREFIX]

    late_ack: bool = Field(
        default=True,
        description="Flag specifying acknowledgment pattern. If set to True "
        "will ack message after execution. Else if will ack the "
        "message before execution.",
    )
    execution_retries: int = Field(
        default=0,
        description="Number of execution retries.",
        ge=0,
    )
    ack_retries: int = Field(
        default=3,
        description="Number of times consumer will retry to ack message.",
        ge=0,
    )
    ack_retry_backoff: int = Field(
        default=1,
        description="Number of seconds between ack message retries.",
        ge=1,
    )
    process_retry_backoff: int = Field(
        default=1,
        description="Number of seconds between process message retries.",
        ge=1,
    )


class Executor(ABC):
    """Abstract executor interface."""

    @abstractmethod
    async def execute(self, payload: MessagePayload) -> None:
        """Executor main method.

        Responsible with executing the payload and adding implementing
        additional layers of error handling and retries:

        - Should decide whether to propagate errors to the queue layer.
        - Should decide whether retries are applicable.

        Args:
            payload: A standard message payload.
        """

    @abstractmethod
    async def shutdown(self) -> None:
        """Graceful executor shutdown (completes work, then exits)."""
        pass

    @abstractmethod
    def max_message_capacity(self) -> int | None:
        """Helper utility for variable batch size consumption.

        Helpful when execution needs to be throttled. The executor
        responds with the number of messages it can pickup for execution
        at the moment.

        Returns:
            Number of messages the executor can handle (None if no limit).
        """


class ConsumerBase(ABC):
    """Async core pipeline shared by polling and push transports.

    The transport decides *how* raw messages arrive. The base class provides:
      - decode/preprocess
      - execute with retries
      - ack with retries (+ critical event on exhaustion)
      - invalid/failing handlers
    """

    def __init__(
        self,
        config: ConsumerRuntimeConfig,
        executor: Executor,
        event_handler: CriticalEventHandler | None = None,
    ) -> None:
        """Consumer base constructor.

        Args:
            config: The consumer config object.
            executor: An executor instance.
            event_handler: A CriticalEventHandler instance (optional).
        """
        self._cfg = config
        self._executor = executor
        self._event_handler = event_handler

    @property
    def config(self) -> ConsumerRuntimeConfig:
        """Implements the 'config' property.

        Returns:
            The consumer config object.
        """
        return self._cfg

    # -------- backend hooks (required) --------

    @abstractmethod
    async def preprocess_message(self, raw_message: Any) -> MessageEnvelope:
        """Decode/Validate broker message and wrap to MessageEnvelope.

        Args:
            raw_message: Broker-specific message.

        Returns:
            A MessageEnvelope object.
        """

    async def process_message(self, message: MessageEnvelope) -> None:
        """Execution of message payload.

        Args:
            message: A MessageEnvelope object.
        """
        await self._executor.execute(payload=message.payload)

    @abstractmethod
    async def acknowledge_message(self, raw_message: Any) -> None:
        """Acks/Deletes message from the queue.

        Args:
            raw_message: The raw broker-specific message.
        """

    @abstractmethod
    async def handle_invalid_message(self, raw_message: Any) -> None:
        """Handles message of invalid format (not suitable for execution).

        Messages that are invalid should still be acked to avoid
        looping execution and system poisoning.

        Args:
            raw_message: The raw broker-specific message.

        """

    async def handle_critical_event(self, event: CriticalEvent) -> None:
        """Critical event handler.

        A critical event indicates that the application behavior has degraded.
        For instance, the inability to serialize the message to a broker-compatible format
        is extremely severe as it most probably repeats across messages meaning no messages
        are published for execution at all.

        A complete implementation of a critical event handling process would be the following:

        1. Log in detailed fashion the event so that it is traceable in the logs.
        2. Publish the event to an event aggregator service so that it is tracked long-term and
        alerts can be generated based on its frequency and severity.
        3. Hold the message in a DLQ to enable re-execution when a bug fix is introduced in the code.

        Args:
            event: A critical event instance.
        """
        logger.exception(
            "Critical event %s caused by the following exception:",
            event.value,
            exc_info=event.exception,
        )
        if self._event_handler:
            asyncio.create_task(self._event_handler.handle(event))

    async def handle_raw_message(self, raw_message: Any) -> None:
        """Handles one message at a time.

        Args:
            raw_message: A raw broker-specific message.
        """
        try:
            msg = await self.preprocess_message(raw_message)
        except MessageDecodeError:
            await self._ack_with_retries(raw_message=raw_message)
            await self.handle_invalid_message(raw_message)
            return

        if self.config.late_ack:
            await self._process_with_retries(msg)
            await self._ack_with_retries(raw_message=msg.raw_message)
        else:
            await self._ack_with_retries(raw_message=msg.raw_message)
            await self._process_with_retries(msg)

    async def _process_with_retries(self, message: MessageEnvelope) -> None:
        attempts = self.config.execution_retries + 1
        last_exc = None

        for attempt in range(attempts):
            try:
                await self.process_message(message)
                return
            except Exception as exc:
                last_exc = exc
                await asyncio.sleep(self.config.process_retry_backoff)

        await self.handle_critical_event(
            CriticalEvent(
                value=CriticalEventType.EXECUTE_FAILED,
                description=f"Execution failed for message {message.id}",
                created_at=datetime.now(tz=timezone.utc),
                exception=last_exc,
            )
        )

    async def _ack_with_retries(self, raw_message: Any) -> None:
        attempts = self.config.ack_retries + 1

        last_exc = None

        for attempt in range(1, attempts + 1):
            try:
                await self.acknowledge_message(raw_message=raw_message)
                return
            except Exception as exc:
                last_exc = exc
                await asyncio.sleep(self.config.ack_retry_backoff)

        await self.handle_critical_event(
            CriticalEvent(
                value=CriticalEventType.ACK_FAILED,
                description="Ack failed for message",
                created_at=datetime.now(tz=timezone.utc),
                exception=last_exc,
            )
        )

    @abstractmethod
    async def run(self) -> None:
        """Starts the consumer task."""
        pass


class PollingConfig(ConsumerRuntimeConfig):
    """Base config class for polling consumers."""

    @staticmethod
    def prefixes() -> list[str]:
        """Implementation of the prefixes config method.

        Returns:
            A list of prefixes to use for env var name resolution.
        """
        return ConsumerRuntimeConfig.prefixes() + [
            ENV_ZENML_CONSUMER_POLLING_PREFIX
        ]

    interval: float = Field(
        default=10.0,
        description="Polling interval in seconds.",
        ge=0.0,
    )
    jitter: float = Field(
        default=0.0,
        description="Jitter(+- random diff) for interval in seconds.",
        ge=0.0,
    )


class PollingConsumer(ConsumerBase, ABC):
    """Polling transport: the consumer fetches batches periodically."""

    def __init__(
        self,
        config: PollingConfig,
        executor: Executor,
        event_handler: CriticalEventHandler | None = None,
    ) -> None:
        """Polling consumer constructor.

        Args:
            config: Polling config instance.
            executor: Executor instance.
            event_handler: A CriticalEventHandler instance (optional).
        """
        super().__init__(
            config, executor=executor, event_handler=event_handler
        )
        self._stopped: bool = False

        # Jitter is computed once at construction (fixed for this consumer instance).
        self._polling_interval = config.interval + random.uniform(
            -config.jitter, config.jitter
        )

    def stop(self) -> None:
        """Signal the run loop to stop."""
        self._stopped = True

    @property
    def polling_interval(self) -> float:
        """Implement the 'polling_interval' property.

        Returns:
            The polling interval value.
        """
        return self._polling_interval

    @abstractmethod
    async def receive_messages(
        self, batch_size: int | None = None
    ) -> list[Any]:
        """Requests to get a batch of messages from the queue.

        Args:
            batch_size: Number of messages to get.

        Returns:
            A list of messages (broker-specific payload).
        """

    @abstractmethod
    def get_messages_batch_size(self) -> int:
        """Getter for the batch size.

        Enables implementations that adapt batch size per iteration
        based on configuration, throttling or other factors.

        Returns:
            The size of the next batch of messages.
        """

    async def poll_once(self) -> None:
        """Receive and handle one batch. Receive failures emit read_failed."""
        try:
            batch_size = self.get_messages_batch_size()
            raw_messages = await self.receive_messages(batch_size=batch_size)
        except Exception as exc:
            await self.handle_critical_event(
                CriticalEvent(
                    value=CriticalEventType.READ_FAILED,
                    description=f"Receive failed: {exc!r}",
                    created_at=datetime.now(tz=timezone.utc),
                    exception=exc,
                )
            )
            return

        for raw in raw_messages:
            await self.handle_raw_message(raw)

    async def run(self) -> None:
        """Run polling loop until stop() is called."""
        try:
            while not self._stopped:
                try:
                    start = asyncio.get_running_loop().time()
                    await self.poll_once()

                    elapsed = asyncio.get_running_loop().time() - start

                    if elapsed > self._polling_interval:
                        logger.debug("Elapsed time for polling: %s", elapsed)
                    remaining = self.polling_interval - elapsed
                    if remaining > 0:
                        await asyncio.sleep(remaining)
                except Exception as exc:
                    logger.exception(
                        "Unhandled exception during polling: %s", exc_info=exc
                    )
        finally:
            self.stop()
            await self._executor.shutdown()
