"""
Helpers for propagating OpenTelemetry trace context through SQS message processing.
"""
import logging
from contextlib import contextmanager
from typing import Any, List, Optional

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

logger = logging.getLogger(__name__)

_propagator = TraceContextTextMapPropagator()


@contextmanager
def sqs_trace_context(traceparent: Optional[str], tenant_id: Optional[str] = None):
    """
    Context manager that restores trace context from a W3C traceparent string.
    Use when processing SQS messages that carry traceparent from upstream requests.

    Args:
        traceparent: W3C traceparent header value (e.g. "00-<trace_id>-<span_id>-01")
        tenant_id: Optional tenant ID (for use with tenant_context if needed)
    """
    token = None
    if traceparent:
        try:
            carrier = {"traceparent": traceparent}
            ctx = _propagator.extract(carrier=carrier)
            token = context.attach(ctx)
        except Exception as e:
            logger.debug("Failed to extract trace context from traceparent: %s", e)

    try:
        yield
    finally:
        if token is not None:
            try:
                context.detach(token)
            except Exception:
                pass


def dispatch_task_with_sqs_trace(
    handler: Any,
    task_args: List[Any],
    traceparent: Optional[str],
    tenant_id: Optional[str] = None,
) -> Any:
    """
    Dispatch a Celery task with trace context from an SQS message.
    Injects traceparent into task headers so worker spans correlate with the upstream trace.

    Args:
        handler: Celery task (e.g. parse_and_create_accounting_codes)
        task_args: Positional arguments for the task
        traceparent: W3C traceparent from the SQS message
        tenant_id: Optional tenant ID

    Returns:
        Result of handler.apply_async()
    """
    headers = {}
    if traceparent:
        headers["traceparent"] = traceparent

    return handler.apply_async(args=task_args, headers=headers)
