"""
Context utilities for managing OpenTelemetry baggage and context propagation.

This module provides helper functions for properly setting and attaching
baggage to the current execution context, ensuring proper propagation
across async boundaries.
"""

from typing import Optional, Dict, Any
from opentelemetry import baggage, context


def set_baggage_context(baggage_items: Dict[str, Any]) -> context.Token:
    """
    Set multiple baggage items and attach them to the current context.
    
    This is the proper way to set baggage that will be propagated
    across async boundaries and HTTP requests.
    
    Args:
        baggage_items: Dictionary of baggage key-value pairs to set
        
    Returns:
        The context token that should be detached when done
        
    Example:
        token = set_baggage_context({
            "user_id": "user@example.com",
            "tenant_id": "tenant-123", 
            "session_id": "session-456",
            "access_token": "bearer-token"
        })
        try:
            # Make calls that need the baggage context
            result = await some_client.make_request()
        finally:
            context.detach(token)
    """
    ctx = context.get_current()
    
    # Set each baggage item
    for key, value in baggage_items.items():
        if value is not None:  # Only set non-None values
            ctx = baggage.set_baggage(key, str(value), ctx)
    
    # Attach the context to make it active
    return context.attach(ctx)


def get_baggage_value(key: str, default: Optional[str] = None) -> Optional[str]:
    """
    Get a baggage value from the current context.
    
    Args:
        key: The baggage key to retrieve
        default: Default value if key not found
        
    Returns:
        The baggage value or default
    """
    current_ctx = context.get_current()
    value = baggage.get_baggage(key, current_ctx)
    return str(value) if value is not None else default


def get_current_baggage() -> Dict[str, str]:
    """
    Get all baggage items from the current context.
    
    Returns:
        Dictionary of all baggage key-value pairs
    """
    current_ctx = context.get_current()
    
    # Common baggage keys we expect
    expected_keys = ["user_id", "tenant_id", "session_id", "access_token", "mission_id"]
    
    baggage_dict = {}
    for key in expected_keys:
        value = baggage.get_baggage(key, current_ctx)
        if value is not None:
            baggage_dict[key] = str(value)
    
    return baggage_dict


def with_baggage_context(baggage_items: Dict[str, Any]):
    """
    Decorator for functions that need baggage context.
    
    Args:
        baggage_items: Dictionary of baggage to set for the function
        
    Example:
        @with_baggage_context({"user_id": "user@example.com", "tenant_id": "tenant-123"})
        async def my_function():
            # Function runs with the baggage context
            pass
    """
    def decorator(func):
        async def wrapper(*args, **kwargs):
            token = set_baggage_context(baggage_items)
            try:
                return await func(*args, **kwargs)
            finally:
                context.detach(token)
        return wrapper
    return decorator
