"""Cross-process trace propagation for multi-agent tracing."""

from collections.abc import Iterator
from contextlib import contextmanager

from opentelemetry import context as otel_context
from opentelemetry import propagate
from opentelemetry import trace as trace_api
from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags

from plyra_trace.models import SpanLink, TraceContext


def inject_context(
    carrier: dict,
    agent_name: str | None = None,
    handoff_reason: str | None = None,
    handoff_protocol: str = "http",
) -> dict:
    """
    Inject current trace context into a carrier dict (HTTP headers).
    Adds W3C traceparent + plyra-specific agent headers.

    Args:
        carrier: Dictionary to inject context into (e.g., HTTP headers)
        agent_name: Optional agent name performing the handoff
        handoff_reason: Optional reason for the handoff
        handoff_protocol: Protocol used for handoff (default: "http")

    Returns:
        The carrier dict with injected context

    Example:
        >>> headers = {}
        >>> inject_context(
        ...     headers,
        ...     agent_name="orchestrator",
        ...     handoff_reason="research task"
        ... )
        >>> response = httpx.post(url, json=payload, headers=headers)
    """
    # Inject W3C trace context using OpenTelemetry's standard propagator
    propagate.inject(carrier)

    # Add plyra-specific headers for agent metadata
    if agent_name:
        carrier["x-plyra-agent-from"] = agent_name
    if handoff_reason:
        carrier["x-plyra-handoff-reason"] = handoff_reason
    carrier["x-plyra-handoff-protocol"] = handoff_protocol

    return carrier


def extract_context(carrier: dict) -> TraceContext:
    """
    Extract trace context from a carrier dict (HTTP headers).
    Returns a TraceContext that can be passed to continue_trace().

    Args:
        carrier: Dictionary containing trace context (e.g., HTTP headers)

    Returns:
        TraceContext object

    Example:
        >>> ctx = extract_context(request.headers)
        >>> with continue_trace(ctx):
        ...     result = do_work(request.body)
    """
    # Extract context using OpenTelemetry's standard propagator
    ctx = propagate.extract(carrier)

    # Get span context from the extracted context
    span = trace_api.get_current_span(ctx)
    span_context = span.get_span_context()

    # Extract plyra-specific headers
    agent_from = carrier.get("x-plyra-agent-from")
    handoff_reason = carrier.get("x-plyra-handoff-reason")
    handoff_protocol = carrier.get("x-plyra-handoff-protocol", "http")

    # The traceparent header format is: version-trace_id-span_id-trace_flags
    # e.g., "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
    # We need to extract trace_id and span_id in hex format

    trace_id_hex = format(span_context.trace_id, "032x")
    span_id_hex = format(span_context.span_id, "016x")

    return TraceContext(
        trace_id=trace_id_hex,
        span_id=span_id_hex,
        trace_flags=span_context.trace_flags,
        agent_from=agent_from,
        handoff_reason=handoff_reason,
        handoff_protocol=handoff_protocol,
        raw_carrier=dict(carrier),
    )


@contextmanager
def continue_trace(ctx: TraceContext) -> Iterator[None]:
    """
    Continue an existing trace from a TraceContext.
    Creates a child span in the originating trace.

    Args:
        ctx: TraceContext extracted from carrier

    Example:
        >>> ctx = extract_context(request.headers)
        >>> with continue_trace(ctx):
        ...     result = run_agent(payload)
    """
    # Create a SpanContext from the trace context
    span_context = SpanContext(
        trace_id=int(ctx.trace_id, 16),
        span_id=int(ctx.span_id, 16),
        is_remote=True,
        trace_flags=TraceFlags(ctx.trace_flags),
    )

    # Create a non-recording span to represent the parent
    parent_span = NonRecordingSpan(span_context)

    # Set this as the current span
    token = otel_context.attach(trace_api.set_span_in_context(parent_span))

    try:
        yield
    finally:
        otel_context.detach(token)


def create_link(
    trace_id: str,
    span_id: str,
    attributes: dict | None = None,
) -> SpanLink:
    """
    Create a span link to a span in another trace.
    Used for async/queue-based cross-agent correlation.

    Args:
        trace_id: Trace ID in hex format
        span_id: Span ID in hex format
        attributes: Optional attributes for the link

    Returns:
        SpanLink object

    Example:
        >>> link = create_link(
        ...     trace_id=message.metadata["source_trace_id"],
        ...     span_id=message.metadata["source_span_id"],
        ...     attributes={
        ...         "agent.handoff.protocol": "sqs",
        ...         "agent.handoff.from": "orchestrator"
        ...     }
        ... )
        >>> with SpanBuilder("process", links=[link]) as span:
        ...     result = do_work(message.body)
    """
    return SpanLink(
        trace_id=trace_id,
        span_id=span_id,
        attributes=attributes or {},
    )
