"""Utilities for tracing messages."""
import base64
import json

from getstream_connector.utils import extractor

from notifications_delivery.connectors import loggly

logger = loggly.get_current_logger()


def extract_trace_context(message):
    """Extract the trace context and/or correlation id from the message.

    Args:
        message: message to extract trace context from

    Returns:
        tuple(trace_contex, correlation_id): tuple of datadog trace context and correlation id
    """
    try:
        # try to decode base64 encoded message
        decoded_message = base64.b64decode(message.body)
    except (TypeError, ValueError):
        # looks like message is not b64encoded (most probably JSON)
        decoded_message = message.body

    try:
        # try to load JSON message
        decoded_message = json.loads(decoded_message)
    except (TypeError, ValueError) as e:
        logger.warning("Can't decode message", exc_info=e)
        return None, None

    # try to get Datadog trace contex and correlation id
    # that's only one supported source of trace context for now
    trace_context = extractor.extract_trace_context(decoded_message)
    correlation_id = extractor.extract_correlation_id(decoded_message)

    return trace_context, correlation_id