"""
Helpers for propagating OpenTelemetry trace context to Celery tasks.
"""
import logging
from contextlib import contextmanager
from typing import Any, Callable, Optional

from opentelemetry import context
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

logger = logging.getLogger(__name__)

_propagator = TraceContextTextMapPropagator()


def call_task_with_trace_context(task: Callable, *args: Any, **kwargs: Any) -> Any:
    """
    Dispatch a Celery task with the current trace context propagated via headers.
    Ensures worker logs and spans are correlated with the originating request.

    Args:
        task: Celery task (e.g. from @shared_task)
        *args: Positional args for apply_async
        **kwargs: Keyword args for apply_async (e.g. queue="queue_A")

    Returns:
        Result of task.apply_async()
    """
    carrier: dict = {}
    current_ctx = context.get_current()
    _propagator.inject(carrier, current_ctx)

    headers = kwargs.pop("headers", {}) or {}
    headers.update(carrier)

    return task.apply_async(*args, headers=headers, **kwargs)


@contextmanager
def tenant_context(tenant_id: str, set_tenant_func: Optional[Callable[[str], Any]] = None):
    """
    Context manager that sets the current tenant for the duration of the block.
    Used when running code (e.g. in Celery workers) that needs tenant context.

    Args:
        tenant_id: Tenant ID to set
        set_tenant_func: Optional callback to set tenant. If not provided, tries
            to use varicon.tenant_filter.get_and_set_tenant_obj when available.
    """
    set_tenant = set_tenant_func
    if set_tenant is None:
        try:
            from varicon.tenant_filter import get_and_set_tenant_obj
            set_tenant = get_and_set_tenant_obj
        except ImportError:
            logger.debug("varicon.tenant_filter not available, tenant_context is a no-op")
            set_tenant = lambda _: None

    try:
        if set_tenant:
            set_tenant(tenant_id)
        yield
    finally:
        try:
            from varicon.tenant_filter import set_current_tenant
            set_current_tenant(None)
        except ImportError:
            pass
