__init__.py

from __future__ import annotations

import functools
import logging
from collections.abc import Callable
from opentelemetry import trace
from typing import Any

from workflows.recipe.recipe import Recipe
from workflows.recipe.validate import validate_recipe
from workflows.recipe.wrapper import RecipeWrapper

__all__ = [
    "Recipe",
    "RecipeWrapper",
    "validate_recipe",
    "wrap_subscribe",
    "wrap_subscribe_broadcast",
]

logger = logging.getLogger("workflows.recipe")


def _wrap_subscription(
    transport_layer,
    subscription_call,
    channel,
    callback,
    *args,
    mangle_for_receiving: Callable[[Any], Any] | None = None,
    **kwargs,
):
    """Internal method to create an intercepting function for incoming messages
    to interpret recipes. This function is then used to subscribe to a channel
    on the transport layer.
      :param transport_layer: Reference to underlying transport object.
      :param subscription_call: Reference to the subscribing function of the
                                transport layer.
      :param channel:  Channel name to subscribe to.
      :param callback: Real function to be called when messages are received.
                       The callback will pass three arguments,
                       a RecipeWrapper object (details below), the header as
                       a dictionary structure, and the message.

      :param allow_non_recipe_messages: Pass on incoming messages that do not
                       include recipe information. In this case the first
                       argument to the callback function will be 'None'.
      :param log_extender: If the recipe contains useful contextual information
                       for log messages, such as a unique ID which can be used
                       to connect all messages originating from the same
                       recipe, then the information will be passed to this
                       function, which must be a context manager factory.
      :return:         Return value of call to subscription_call.
    """

    allow_non_recipe_messages = kwargs.pop("allow_non_recipe_messages", False)
    log_extender = kwargs.pop("log_extender", None)

    @functools.wraps(callback)
    def unwrap_recipe(header, message):
        """This is a helper function unpacking incoming messages when they are
        in a recipe format. Other messages are passed through unmodified.
        :param header:  A dictionary of message headers. If the header contains
                        an entry 'workflows-recipe' then the message is parsed
                        and the embedded recipe information is passed on in a
                        RecipeWrapper object to the target function.
        :param message: Incoming deserialized message object.
        """
        if mangle_for_receiving:
            message = mangle_for_receiving(message)
        if header.get("workflows-recipe") in {True, "True", "true", 1}:
            rw = RecipeWrapper(message=message, transport=transport_layer)
            logger.debug("RecipeWrapper created: %s", rw)

            # Extract recipe_id on the current span
            span = trace.get_current_span()
            dcid = None
            recipe_id = None

            # Extract recipe ID from environment
            if isinstance(message, dict):
                environment = message.get("environment", {})
                if isinstance(environment, dict):
                    recipe_id = environment.get("ID")
            
            if recipe_id:
                span.set_attribute("recipe_id", recipe_id)
                span.add_event("recipe.id_extracted", attributes={"recipe_id": recipe_id})

            # Extract span_id and trace_id for logging
            span_context = span.get_span_context()
            if span_context and span_context.is_valid:
                span_id = format(span_context.span_id, '016x')
                trace_id = format(span_context.trace_id, '032x')

                log_extra = {
                    "span_id": span_id,
                    "trace_id": trace_id,
                }
                if dcid:
                    log_extra["dcid"] = dcid
                if recipe_id:
                    log_extra["recipe_id"] = recipe_id

                logger.info(

                    "Processing recipe message",
                    extra=log_extra
                )
        
            if log_extender and rw.environment and rw.environment.get("ID"):
                with log_extender("recipe_ID", rw.environment["ID"]):
                    return callback(rw, header, message.get("payload"))
            return callback(rw, header, message.get("payload"))
        if allow_non_recipe_messages:
            return callback(None, header, message)
        #   self.log.warning('Discarding non-recipe message:\n' + \
        #                    "First 1000 characters of header:\n%s\n" + \
        #                    "First 1000 characters of message:\n%s",
        #                    str(header)[:1000], str(message)[:1000])
        logger.error(
            "The input to this service is not a wrapped recipe. "
            "Unable to process incoming message."
        )
        transport_layer.nack(header)

    if mangle_for_receiving:
        kwargs = {**kwargs, "disable_mangling": True}
    return subscription_call(channel, unwrap_recipe, *args, **kwargs)


def wrap_subscribe(
    transport_layer,
    channel,
    callback,
    *args,
    mangle_for_receiving: Callable[[Any], Any] | None = None,
    **kwargs,
):
    """Listen to a queue on the transport layer, similar to the subscribe call in
    transport/common_transport.py. Intercept all incoming messages and parse
    for recipe information.
    See common_transport.subscribe for possible additional keyword arguments.
      :param transport_layer: Reference to underlying transport object.
      :param channel:  Queue name to subscribe to.
      :param callback: Function to be called when messages are received.
                       The callback will pass three arguments,
                       a RecipeWrapper object (details below), the header as
                       a dictionary structure, and the message.
      :return: A unique subscription ID
    """

    return _wrap_subscription(
        transport_layer,
        transport_layer.subscribe,
        channel,
        callback,
        *args,
        mangle_for_receiving=mangle_for_receiving,
        **kwargs,
    )


def wrap_subscribe_broadcast(
    transport_layer,
    channel,
    callback,
    *args,
    mangle_for_receiving: Callable[[Any], Any] | None = None,
    **kwargs,
):
    """Listen to a topic on the transport layer, similar to the
    subscribe_broadcast call in transport/common_transport.py. Intercept all
    incoming messages and parse for recipe information.
    See common_transport.subscribe_broadcast for possible arguments.
      :param transport_layer: Reference to underlying transport object.
      :param channel:  Topic name to subscribe to.
      :param callback: Function to be called when messages are received.
                       The callback will pass three arguments,
                       a RecipeWrapper object (details below), the header as
                       a dictionary structure, and the message.
      :return: A unique subscription ID
    """

    return _wrap_subscription(
        transport_layer,
        transport_layer.subscribe_broadcast,
        channel,
        callback,
        *args,
        mangle_for_receiving=mangle_for_receiving,
        **kwargs,
    )