"""
Context management for Noveum Trace SDK.

This module handles trace context propagation across function calls,
async operations, and thread boundaries.
"""

import contextvars
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Optional

from noveum_trace.core.span import Span, SpanStatus
from noveum_trace.core.trace import Trace


@dataclass
class TraceContext:
    """
    Container for trace context information.

    This class holds the current trace and span context that
    propagates across function calls and async operations.
    """

    trace: Optional[Trace] = None
    span: Optional[Span] = None
    attributes: dict[str, Any] = field(default_factory=dict)
    project_id: Optional[str] = None
    organization_id: Optional[str] = None
    environment: Optional[str] = None

    def __post_init__(self) -> None:
        """Initialize attributes if None."""
        if self.attributes is None:
            self.attributes = {}


# Context variables for trace propagation
_trace_context: contextvars.ContextVar[Optional[TraceContext]] = contextvars.ContextVar(
    "noveum_trace_context", default=None
)


def get_current_context() -> TraceContext:
    """
    Get the current trace context.

    Returns:
        Current TraceContext instance
    """
    context = _trace_context.get()
    if context is None:
        context = TraceContext()
        _trace_context.set(context)
    return context


def get_current_trace() -> Optional[Trace]:
    """
    Get the current trace from context.

    Returns:
        Current Trace instance or None if no trace is active
    """
    context = get_current_context()
    return context.trace


def get_current_span() -> Optional[Span]:
    """
    Get the current span from context.

    Returns:
        Current Span instance or None if no span is active
    """
    context = get_current_context()
    return context.span


def set_current_trace(trace: Optional[Trace]) -> None:
    """
    Set the current trace in context.

    Args:
        trace: Trace instance to set as current
    """
    context = get_current_context()
    context.trace = trace
    _trace_context.set(context)


def set_current_span(span: Optional[Span]) -> None:
    """
    Set the current span in context.

    Args:
        span: Span instance to set as current
    """
    context = get_current_context()
    context.span = span
    _trace_context.set(context)


def set_current_context(context: TraceContext) -> None:
    """
    Set the current trace context.

    Args:
        context: TraceContext instance to set as current
    """
    _trace_context.set(context)


def set_context_attribute(key: str, value: Any) -> None:
    """
    Set an attribute in the current context.

    Args:
        key: Attribute key
        value: Attribute value
    """
    context = get_current_context()
    context.attributes[key] = value
    _trace_context.set(context)


def get_context_attribute(key: str, default: Any = None) -> Any:
    """
    Get an attribute from the current context.

    Args:
        key: Attribute key
        default: Default value if key not found

    Returns:
        Attribute value or default
    """
    context = get_current_context()
    return context.attributes.get(key, default)


@contextmanager
def trace_context(
    trace: Optional[Trace] = None, span: Optional[Span] = None, **attributes: Any
) -> Generator[TraceContext, None, None]:
    """
    Context manager for setting trace context.

    Args:
        trace: Trace to set as current
        span: Span to set as current
        **attributes: Additional context attributes

    Yields:
        TraceContext instance

    Example:
        >>> with trace_context(user_id="123", session_id="abc"):
        ...     # Code here runs with the specified context
        ...     result = some_function()
    """
    # Get current context
    current_context = get_current_context()

    # Create new context with provided values
    new_context = TraceContext(
        trace=trace or current_context.trace,
        span=span or current_context.span,
        attributes={**current_context.attributes, **attributes},
    )

    # Set new context
    token = _trace_context.set(new_context)

    try:
        yield new_context
    finally:
        # Restore previous context
        _trace_context.reset(token)


@contextmanager
def span_context(span: Span) -> Generator[Span, None, None]:
    """
    Context manager for setting a span as current.

    Args:
        span: Span to set as current

    Yields:
        The span instance

    Example:
        >>> span = trace.create_span("operation")
        >>> with span_context(span):
        ...     # Code here runs with span as current
        ...     result = some_function()
    """
    # Get current context
    current_context = get_current_context()

    # Create new context with the span
    new_context = TraceContext(
        trace=current_context.trace,
        span=span,
        attributes=current_context.attributes.copy(),
    )

    # Set new context
    token = _trace_context.set(new_context)

    try:
        yield span
    finally:
        # Restore previous context
        _trace_context.reset(token)


def copy_context() -> TraceContext:
    """
    Create a copy of the current trace context.

    Returns:
        Copy of current TraceContext
    """
    current = get_current_context()
    return TraceContext(
        trace=current.trace, span=current.span, attributes=current.attributes.copy()
    )


def clear_context() -> None:
    """Clear the current trace context."""
    _trace_context.set(TraceContext())


def propagate_context_to_thread(context: Optional[TraceContext] = None) -> TraceContext:
    """
    Propagate trace context to a new thread.

    Args:
        context: Context to propagate (defaults to current context)

    Returns:
        Context to use in the new thread

    Example:
        >>> import threading
        >>>
        >>> def worker():
        ...     # Restore context in thread
        ...     context = propagate_context_to_thread()
        ...     _trace_context.set(context)
        ...     # Now the thread has access to the trace context
        ...
        >>> thread = threading.Thread(target=worker)
        >>> thread.start()
    """
    if context is None:
        context = get_current_context()

    return TraceContext(
        trace=context.trace, span=context.span, attributes=context.attributes.copy()
    )


class ContextualSpan:
    """
    A span that automatically manages context.

    This class wraps a Span and automatically sets it as the current
    span when used as a context manager.
    """

    def __init__(self, span: Span):
        """
        Initialize contextual span.

        Args:
            span: Span to wrap
        """
        self.span = span
        self._previous_span: Optional[Span] = None

    def __enter__(self) -> Span:
        """Context manager entry."""
        self._previous_span = get_current_span()
        set_current_span(self.span)
        return self.span

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Context manager exit."""
        if exc_type is not None:
            self.span.record_exception(exc_val)
            self.span.set_status(SpanStatus.ERROR, str(exc_val))

        # Use client to finish span for consistency and proper export
        from noveum_trace.core import get_global_client

        client = get_global_client()
        if client:
            client.finish_span(self.span)
        else:
            # Fallback to direct span finish if client not available
            self.span.finish()

        set_current_span(self._previous_span)

    async def __aenter__(self) -> Span:
        """Async context manager entry."""
        self._previous_span = get_current_span()
        set_current_span(self.span)
        return self.span

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Async context manager exit."""
        if exc_type is not None:
            self.span.record_exception(exc_val)
            self.span.set_status(SpanStatus.ERROR, str(exc_val))

        # Use client to finish span for consistency and proper export
        from noveum_trace.core import get_global_client

        client = get_global_client()
        if client:
            client.finish_span(self.span)
        else:
            # Fallback to direct span finish if client not available
            self.span.finish()

        set_current_span(self._previous_span)

    def __getattr__(self, name: str) -> Any:
        """Delegate attribute access to the wrapped span."""
        return getattr(self.span, name)


class ContextualTrace:
    """
    A trace that automatically manages context.

    This class wraps a Trace and automatically sets it as the current
    trace when used as a context manager.
    """

    def __init__(self, trace: Trace):
        """
        Initialize contextual trace.

        Args:
            trace: Trace to wrap
        """
        self.trace = trace
        self._previous_trace: Optional[Trace] = None

    def __enter__(self) -> Trace:
        """Context manager entry."""
        self._previous_trace = get_current_trace()
        set_current_trace(self.trace)
        return self.trace

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Context manager exit."""
        if exc_type is not None:
            self.trace.set_status(SpanStatus.ERROR, str(exc_val))

        # Use client to finish trace for proper export and context management
        from noveum_trace.core import get_global_client

        client = get_global_client()
        if client:
            client.finish_trace(self.trace)
        else:
            # Fallback to direct trace finish if client not available
            self.trace.finish()

        set_current_trace(self._previous_trace)

    async def __aenter__(self) -> Trace:
        """Async context manager entry."""
        self._previous_trace = get_current_trace()
        set_current_trace(self.trace)
        return self.trace

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Async context manager exit."""
        if exc_type is not None:
            self.trace.set_status(SpanStatus.ERROR, str(exc_val))

        # Use client to finish trace for proper export and context management
        from noveum_trace.core import get_global_client

        client = get_global_client()
        if client:
            client.finish_trace(self.trace)
        else:
            # Fallback to direct trace finish if client not available
            self.trace.finish()

        set_current_trace(self._previous_trace)

    def __getattr__(self, name: str) -> Any:
        """Delegate attribute access to the wrapped trace."""
        return getattr(self.trace, name)


def attach_context_to_span(span: Span) -> None:
    """
    Attach current context attributes to a span.

    Args:
        span: Span to attach context to
    """
    context = get_current_context()

    # Add context attributes to span
    # Defensive: ensure attributes is a dict (should always be, but be safe)
    if not isinstance(context.attributes, dict):
        span.set_attributes({})
        return

    context_attributes = {}
    for key, value in context.attributes.items():
        context_attributes[f"context.{key}"] = value

    if context_attributes:
        span.set_attributes(context_attributes)
    else:
        # Call set_attributes with empty dict to match test expectations
        span.set_attributes({})


def inherit_context_attributes(span: Span, parent_span: Optional[Span] = None) -> None:
    """
    Inherit context attributes from parent span.

    Args:
        span: Span to inherit attributes to
        parent_span: Parent span to inherit from (defaults to current span)
    """
    if parent_span is None:
        parent_span = get_current_span()

    if parent_span is None:
        return

    # Inherit specific attributes from parent
    inheritable_attributes = [
        "user_id",
        "session_id",
        "request_id",
        "agent_id",
        "workflow_id",
    ]

    for attr in inheritable_attributes:
        context_key = f"context.{attr}"
        if context_key in parent_span.attributes:
            span.set_attribute(context_key, parent_span.attributes[context_key])