"""
SQS trace context helpers for distributed tracing across SQS queues.

Provides functions to:
- Inject trace context (traceparent) when sending SQS messages
- Extract trace context from received SQS messages and dispatch Celery tasks

Usage (Sender side - Integration Service):
    from varicon_observability.sqs_helpers import get_current_traceparent, set_task_traceparent

    # In Celery task (sync context, before calling async code):
    set_task_traceparent(get_current_traceparent())

    # In async handler, Event model auto-calls get_current_traceparent()
    event = Event(..., traceparent=get_current_traceparent())

Usage (Receiver side - Backend Varicon):
    from varicon_observability.sqs_helpers import dispatch_task_with_sqs_trace

    traceparent = body.get("traceparent")
    tenant_id = body.get("tenant_id")
    dispatch_task_with_sqs_trace(my_celery_task, [arg1, arg2], traceparent, tenant_id)
"""

import logging
import threading
import uuid
from contextlib import contextmanager
from typing import Optional, Any, List

logger = logging.getLogger(__name__)

# Thread-local storage for task-level traceparent.
# Celery prefork workers run one task at a time per process, so thread-local is safe.
# This bridges the sync→async gap: set in sync Celery task code, read in async handlers.
_task_trace = threading.local()


def set_task_traceparent(traceparent: Optional[str] = None):
    """
    Store a traceparent for the current Celery task.

    Call this in the sync Celery task function BEFORE run_until_complete(),
    so that async code can access the traceparent via get_current_traceparent().

    If traceparent is None, generates a new one automatically.

    Args:
        traceparent: W3C traceparent string, or None to auto-generate.

    Example:
        @celery_app.task(queue="worker_integration")
        def sync_suppliers_task(pk, tenant_id):
            set_task_traceparent(get_current_traceparent())
            loop = asyncio.get_event_loop()
            return loop.run_until_complete(async_sync_suppliers(pk, tenant_id))
    """
    if not traceparent:
        # Try to capture the current OTEL span before falling back to random
        traceparent = get_current_traceparent()
    if not traceparent:
        trace_id = uuid.uuid4().hex
        span_id = uuid.uuid4().hex[:16]
        traceparent = f"00-{trace_id}-{span_id}-01"
        logger.info(f"[SQS_TRACE] Generated new traceparent: {traceparent}")
    _task_trace.traceparent = traceparent


def get_current_traceparent() -> Optional[str]:
    """
    Get the current trace context as a W3C traceparent string.

    Uses a 2-tier fallback:
      1. Active OTEL span (best — fully correlated with the request trace)
      2. Task-level traceparent (set via set_task_traceparent in Celery task)

    Returns the traceparent string, or None if neither source is available.
    Safe to call even if OpenTelemetry is not installed.

    Returns:
        str or None: e.g. "00-abc123def456...-1111111111111111-01"
    """
    # 1. Try active OTEL span
    try:
        from opentelemetry import trace

        span = trace.get_current_span()
        if span and span.get_span_context().is_valid:
            ctx = span.get_span_context()
            trace_id = format(ctx.trace_id, "032x")
            span_id = format(ctx.span_id, "016x")
            trace_flags = format(ctx.trace_flags, "02x")
            return f"00-{trace_id}-{span_id}-{trace_flags}"
    except Exception:
        pass

    # 2. Try task-level storage (set by Celery task before async call)
    tp = getattr(_task_trace, "traceparent", None)
    if tp:
        return tp

    return None


def dispatch_task_with_sqs_trace(task, args: List[Any], traceparent: Optional[str] = None, tenant_id: Optional[str] = None):
    """
    Dispatch a Celery task with trace context restored from an SQS message.

    If traceparent is provided and OpenTelemetry is available, this function:
    1. Restores the trace context from the traceparent string
    2. Sets the tenant_id in the observability context
    3. Dispatches the task using call_task_with_trace_context (propagates trace to Celery worker)
    4. Cleans up the context after dispatch

    If traceparent is None or OpenTelemetry is not available, falls back to plain .apply_async().

    Args:
        task: The Celery task to dispatch
        args: List of arguments to pass to the task
        traceparent: W3C traceparent string from the SQS message body (e.g. body.get("traceparent"))
        tenant_id: Tenant ID from the SQS message body

    Example:
        from varicon_observability.sqs_helpers import dispatch_task_with_sqs_trace

        # In sqs_worker.py after parsing the message body:
        traceparent = body.get("traceparent")
        tenant_id = body.get("tenant_id")
        dispatch_task_with_sqs_trace(
            parse_and_update_accounting_codes,
            [payload, tenant_id, source, user],
            traceparent,
            tenant_id,
        )
    """
    if not traceparent:
        task.apply_async(args=args)
        return

    try:
        from opentelemetry import context as otel_context
        from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
        from .celery_helpers import call_task_with_trace_context, tenant_context
    except ImportError:
        # OpenTelemetry not available, fall back to plain dispatch
        task.apply_async(args=args)
        return

    token = None
    try:
        # Normalise bare trace_id (32 hex chars) to full W3C traceparent
        if traceparent and not traceparent.startswith("00-"):
            traceparent = f"00-{traceparent}-{'0' * 16}-01"

        # Restore the trace context from the traceparent string
        propagator = TraceContextTextMapPropagator()
        carrier = {"traceparent": traceparent}
        parent_context = propagator.extract(carrier)

        # Set tenant_id in the observability context
        if tenant_id:
            tenant_context.set(str(tenant_id))

        # Attach the extracted context so it becomes the current context
        token = otel_context.attach(parent_context)
        logger.info(
            f"[SQS_TRACE] Restored trace context: traceparent={traceparent}, tenant_id={tenant_id}"
        )

        # Dispatch task with trace propagation to Celery worker
        call_task_with_trace_context(task, *args)
        logger.info(f"[SQS_TRACE] Dispatched {task.name} with trace context")

    except Exception as e:
        logger.warning(f"[SQS_TRACE] Failed to dispatch with trace context: {e}, falling back to plain dispatch")
        task.apply_async(args=args)

    finally:
        # Clean up: detach context and reset tenant
        if token is not None:
            otel_context.detach(token)
        if tenant_id:
            try:
                tenant_context.set("unknown-tenant")
            except Exception:
                pass


@contextmanager
def sqs_trace_context(traceparent: Optional[str] = None, tenant_id: Optional[str] = None):
    """
    Context manager that restores OTel trace context from a traceparent string.

    All logs emitted inside the `with` block will carry the correct trace_id/span_id,
    making them visible and correlated in Signoz.

    If traceparent is None, a **new root span** is created so that all logs
    inside the block still get a trace_id for correlation in Signoz.

    Safe to use even if OpenTelemetry is not installed — becomes a no-op.

    Usage (in sqs_worker.py):
        from varicon_observability.sqs_helpers import sqs_trace_context

        traceparent = body.get("traceparent")
        tenant_id = body.get("tenant_id")
        with sqs_trace_context(traceparent, tenant_id):
            LOGGER.info("this log will have trace_id in Signoz")
            ...
    """
    try:
        from opentelemetry import trace, context as otel_context
        from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
        from .celery_helpers import tenant_context as _tenant_context
    except ImportError:
        yield
        return

    span = None
    ctx_token = None
    _tenant_ctx = _tenant_context
    try:
        tracer = trace.get_tracer("sqs_worker")

        if traceparent:
            # Normalise bare trace_id (32 hex chars) to full W3C traceparent
            if not traceparent.startswith("00-"):
                traceparent = f"00-{traceparent}-{'0' * 16}-01"

            propagator = TraceContextTextMapPropagator()
            parent_context = propagator.extract({"traceparent": traceparent})

            span = tracer.start_span(name="sqs_worker.process_message", context=parent_context)
            span_context = trace.set_span_in_context(span, parent_context)
        else:
            # No incoming trace — create a new root span so all logs
            # inside this block get a trace_id for Signoz correlation.
            span = tracer.start_span(name="sqs_worker.process_message")
            span_context = trace.set_span_in_context(span)

        ctx_token = otel_context.attach(span_context)

        # Verify span is now current — aids debugging empty trace_id issues
        current_span = trace.get_current_span()
        if current_span and current_span.get_span_context().is_valid:
            _ctx = current_span.get_span_context()
            logger.info(
                f"[SQS_TRACE] Context attached: trace_id={format(_ctx.trace_id, '032x')}, "
                f"span_id={format(_ctx.span_id, '016x')}"
            )
        else:
            logger.warning("[SQS_TRACE] Span not current after attach — trace_id will be empty in logs")

        if tenant_id:
            _tenant_ctx.set(str(tenant_id))
    except Exception as e:
        logger.warning(f"[SQS_TRACE] Failed to set up trace context: {e}")

    try:
        yield
    finally:
        if span is not None:
            span.end()
        if ctx_token is not None:
            try:
                otel_context.detach(ctx_token)
            except Exception:
                pass
        try:
            _tenant_ctx.set("unknown-tenant")
        except Exception:
            pass
