from ddtrace.trace import tracer
from ddtrace._trace.context import Context
from ddtrace.propagation.http import HTTPPropagator

import json
from functools import wraps
from typing import Optional
from ddtrace.trace import tracer
from ddtrace.sampler import DatadogSampler
from ddtrace.sampling_rule import SamplingRule


# env_name is optional and used for ddog sampling for prod, by default set it to dev, which commonly means no sampling
def distributed_trace(span_name: str, service_name: str, env_name: str = "dev"):
    """
    A decorator for generating tracing spans around the execution of the decorated function.

    This function creates a decorator that, when applied to another function, wraps the execution
    of that function in a tracing span. This is useful for monitoring and debugging the performance
    and behavior of the decorated function in a distributed system.

    Args:
        span_name (str): The name of the span that will be created for tracing.
        service_name (str): The name of the service under which the span will be categorized.

    Returns:
        A decorator function that takes a function and returns a wrapped version of that function
        with tracing enabled.
    """

    def decorator(func):
        @wraps(func)
        def wrapper(self, *args, parent_context: Optional[str] = None, **kwargs):
            try:
                if parent_context:
                    parent_context = HTTPPropagator.extract(json.loads(parent_context))
            except Exception as e:
                print("Error extracting parent context:", e)
                parent_context = None

            if (env_name == "prod") or (env_name == "msft"):
                tracer.configure(
                    sampler=DatadogSampler(
                        rules=[
                            SamplingRule(sample_rate=0.0001),
                        ]
                    )
                )

            with tracer.start_span(
                span_name,
                service=service_name,
                child_of=parent_context,
                activate=True,
            ):
                return func(self, *args, **kwargs)

        return wrapper

    return decorator


def serialize_context(context: Optional[Context] = None, http=True) -> str:
    if context is None:
        context = tracer.current_trace_context()
    if http:
        headers = {}
        HTTPPropagator.inject(context, headers)
        return json.dumps(headers)
    else:
        # This is deprecated, for backwards compatibility
        return json.dumps(
            {
                "trace_id": context.trace_id,
                "span_id": context.span_id,
            }
        )