import functools
from typing import Any, Callable, Dict, List, Tuple, Union
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry import context

def with_trace_context(func: Callable):
    @functools.wraps(func)
    async def wrapper(*args, **kwargs) -> Any:
        carrier = kwargs.get('message_headers')
        if not carrier and len(args) >= 3:
            carrier = args[2]
        
        if not carrier or not isinstance(carrier, dict):
            return await func(*args, **kwargs)

        propagator = TraceContextTextMapPropagator()
        ctx = propagator.extract(carrier=carrier)

        token = context.attach(ctx)
        try:
            return await func(*args, **kwargs)
        finally:
            context.detach(token)

    return wrapper