"""Context-aware span processor for injecting Flowlines attributes."""

from __future__ import annotations

from contextvars import ContextVar
from typing import TYPE_CHECKING

from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor

if TYPE_CHECKING:
    from contextvars import Token

    from opentelemetry.context import Context

_flowlines_attrs: ContextVar[dict[str, str] | None] = ContextVar(
    "_flowlines_attrs",
    default=None,
)


class FlowlinesSpanProcessor(SpanProcessor):
    """Span processor that stamps Flowlines context attributes onto spans.

    Wraps an inner ``SpanProcessor`` (typically ``BatchSpanProcessor``),
    injecting any active ``flowlines.*`` attributes in ``on_start`` before
    delegating all lifecycle calls to the inner processor.
    """

    def __init__(self, inner: SpanProcessor) -> None:
        self._inner = inner

    def on_start(
        self,
        span: Span,
        parent_context: Context | None = None,
    ) -> None:
        attrs = _flowlines_attrs.get()
        if attrs is not None:
            for key, value in attrs.items():
                span.set_attribute(key, value)
        self._inner.on_start(span, parent_context)

    def on_end(self, span: ReadableSpan) -> None:
        self._inner.on_end(span)

    def shutdown(self) -> None:
        self._inner.shutdown()

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return self._inner.force_flush(timeout_millis)


def set_context(
    *,
    user_id: str | None = None,
    session_id: str | None = None,
    agent_id: str | None = None,
) -> Token[dict[str, str] | None]:
    """Set Flowlines context attributes on the current context.

    Returns a token that can be passed to ``clear_context`` to restore
    the previous state.
    """
    attrs: dict[str, str] = {}
    if user_id is not None:
        attrs["flowlines.user_id"] = user_id
    if session_id is not None:
        attrs["flowlines.session_id"] = session_id
    if agent_id is not None:
        attrs["flowlines.agent_id"] = agent_id
    return _flowlines_attrs.set(attrs)


def clear_context(token: Token[dict[str, str] | None]) -> None:
    """Restore the previous Flowlines context using the given token."""
    _flowlines_attrs.reset(token)
