"""DataDog wrapper functions."""

import logging
from typing import Any, Collection

from ddtrace import config as dd_config
from ddtrace import tracer
from ddtrace.internal.utils.http import redact_url

from .. import config


logger = logging.getLogger(__name__)


__all__ = ["wrap_methods", "set_tags"]


def wrap_methods(
    obj: Any,
    methods: Collection[str],
    name_prefix: str = "stream",
    service_name: str = config.SERVICE_NAME,
) -> None:
    """Wrap specified methods with datadog trace.

    :param obj: object to get methods from
    :param methods: methods to wrap
    :param name_prefix: prefix to add to DataDog traces
    :param service_name: service name for DataDog traces
    """
    for method_name in methods:
        try:
            method = getattr(obj, method_name)
            wrapper = tracer.wrap(f"{name_prefix}.{method_name}", service=service_name)
            setattr(obj, method_name, wrapper(method))
        except BaseException as e:
            logger.warning(
                f"An error occurred while patching {obj.__class__.__name__}.{method_name}",
                exc_info=e,
            )


def set_tags(status_code: int, url: str | None, method: str | None, ratelimit_info: dict[str, Any]) -> None:
    """Set datadog trace tags."""
    if not tracer.enabled:
        return

    try:
        span = tracer.current_span()
        if not span:
            return

        # Set http tags
        span.set_tag("http.status_code", str(status_code))

        if url is not None:
            redacted_url = redact_url(url, dd_config._obfuscation_query_string_pattern)
            redacted_url = redacted_url.decode("utf-8") if isinstance(redacted_url, bytes) else redacted_url
            span.set_tag("http.url", redacted_url)

        if method is not None:
            span.set_tag("http.method", method.upper())

        # Set app tags
        for key, value in ratelimit_info.items():
            span.set_tag(f"stream.{key}", value)

    except BaseException as e:
        logger.warning("An error occurred while adding tags.", exc_info=e)