"""Context managers for session, user, metadata, and tags propagation."""

import json
from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any

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

from plyra_trace.semconv import SpanAttributes

# Context variables for propagating values to child spans
_session_id_var: ContextVar[str | None] = ContextVar("plyra_session_id", default=None)
_user_id_var: ContextVar[str | None] = ContextVar("plyra_user_id", default=None)
_metadata_var: ContextVar[dict[str, Any] | None] = ContextVar("plyra_metadata", default=None)
_tags_var: ContextVar[list[str] | None] = ContextVar("plyra_tags", default=None)


class ContextSpanProcessor(SpanProcessor):
    """
    Span processor that reads from context variables and sets attributes on new spans.
    This is registered automatically when plyra-trace is initialized.
    """

    def on_start(self, span: Span, parent_context: Context | None = None) -> None:
        """Called when a span is started. Enrich with context attributes."""
        # Get values from context variables
        session_id = _session_id_var.get()
        user_id = _user_id_var.get()
        metadata = _metadata_var.get()
        tags = _tags_var.get()

        # Set attributes on the span
        if session_id is not None:
            span.set_attribute(SpanAttributes.SESSION_ID, session_id)

        if user_id is not None:
            span.set_attribute(SpanAttributes.USER_ID, user_id)

        if metadata is not None:
            span.set_attribute(SpanAttributes.METADATA, json.dumps(metadata))

        if tags is not None:
            span.set_attribute(SpanAttributes.TAG_TAGS, json.dumps(tags))

    def on_end(self, span: ReadableSpan) -> None:
        """Called when a span ends."""
        pass

    def shutdown(self) -> None:
        """Called when the span processor is shutdown."""
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush spans."""
        return True


@contextmanager
def session(session_id: str) -> Iterator[None]:
    """
    Attach session.id to all spans within the block.

    Args:
        session_id: Session identifier

    Example:
        >>> with plyra_trace.session("session-abc-123"):
        ...     result = run_agent(query)
    """
    token = _session_id_var.set(session_id)
    try:
        yield
    finally:
        _session_id_var.reset(token)


@contextmanager
def user(user_id: str) -> Iterator[None]:
    """
    Attach user.id to all spans within the block.

    Args:
        user_id: User identifier

    Example:
        >>> with plyra_trace.user("user-456"):
        ...     result = run_agent(query)
    """
    token = _user_id_var.set(user_id)
    try:
        yield
    finally:
        _user_id_var.reset(token)


@contextmanager
def metadata(data: dict[str, Any]) -> Iterator[None]:
    """
    Attach arbitrary metadata (JSON-serialized) to all spans within the block.

    Args:
        data: Metadata dictionary

    Example:
        >>> with plyra_trace.metadata(
        ...     {"customer_tier": "enterprise", "region": "eu-west"}
        ... ):
        ...     result = run_agent(query)
    """
    # Merge with existing metadata if present
    existing = _metadata_var.get()
    if existing is not None:
        merged = {**existing, **data}
    else:
        merged = data

    token = _metadata_var.set(merged)
    try:
        yield
    finally:
        _metadata_var.reset(token)


@contextmanager
def tags(tag_list: list[str]) -> Iterator[None]:
    """
    Attach tags to all spans within the block.

    Args:
        tag_list: List of tag strings

    Example:
        >>> with plyra_trace.tags(["production", "high-priority"]):
        ...     result = run_agent(query)
    """
    # Merge with existing tags if present
    existing = _tags_var.get()
    if existing is not None:
        merged = list(set(existing + tag_list))  # Deduplicate
    else:
        merged = tag_list

    token = _tags_var.set(merged)
    try:
        yield
    finally:
        _tags_var.reset(token)
