"""Functions to extract data from 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__ = ["extract_trace_context", "extract_correlation_id"]


logger = logging.getLogger(__name__)


def _get_activity_data(message: list[dict[str, Any]]) -> dict[str, Any] | None:
    """Try to get activity data from the message.

    :param message: Decoded GetStream message to extract data from
    :return: Activity data
    """
    try:
        return message[0]["new"][0]
    except (IndexError, KeyError) as e:
        logger.warning("Can't get activity data.", exc_info=e)

    return None


def extract_trace_context(message: list[dict[str, Any]]) -> Union["Context", None]:
    """Extract datadog trace context from the message.

    :param message: Decoded GetStream message to extract data from
    :return: DataDog trace context or None
    """
    data = _get_activity_data(message)

    if not data or config.TRACE_CONTEXT_KEY not in data:
        logger.warning("Can't extract trace context from activity.")
        return None

    return HTTPPropagator.extract(data.pop(config.TRACE_CONTEXT_KEY))


def extract_correlation_id(message: list[dict[str, Any]]) -> str | None:
    """Extract correlation id from the message.

    :param message: Decoded Getstream message to extract data from
    :return: correlation id or None
    """
    data = _get_activity_data(message)

    if not data or config.CORRELATION_ID_KEY not in data:
        logger.warning("Can't extract correlation id from activity.")
        return None

    return data.pop(config.CORRELATION_ID_KEY)