"""Functions to inject data to the GetStream messages."""

import logging
from typing import TYPE_CHECKING, Any, Union

from ddtrace.propagation.http import HTTPPropagator

from .. import config


if TYPE_CHECKING:
    from ddtrace.context import Context


__all__ = ["inject_trace_context", "inject_correlation_id"]


logger = logging.getLogger(__name__)


def inject_trace_context(activity: dict[str, Any], context: Union["Context", None]) -> None:
    """Inject provided tracing context to getstream message for distributed tracing.

    :param activity: activity in which the correlation id must be injected
    :param context: DataDog trace context to inject
    """
    if not activity or not context:
        return

    if config.TRACE_CONTEXT_KEY in activity:
        logger.info("Trace context already exists in activity.")
        return

    try:
        trace_context_dict: dict[str, str] = {}
        HTTPPropagator.inject(context, trace_context_dict)
        activity[config.TRACE_CONTEXT_KEY] = trace_context_dict
    except BaseException as e:
        logger.warning("Can't insert trace context to activity.", exc_info=e)


def inject_correlation_id(activity: dict[str, Any], correlation_id: str | None) -> None:
    """Inject correlation id to getstream message.

    :param activity: activity in which the correlation id must be injected
    :param correlation_id: correlation id to be injected
    :return:
    """
    if not activity or not correlation_id:
        return

    if config.CORRELATION_ID_KEY in activity:
        logger.info("Correlation id already exists in activity.")
        return

    try:
        activity[config.CORRELATION_ID_KEY] = correlation_id
    except BaseException as e:
        logger.warning("Can't insert correlation id to activity.", exc_info=e)