"""AgentBill Rich Context Model (P3-C)

Provides scoped context management for billing attribution,
workflow correlation, and metadata propagation.

[CONTEXT_UPGRADE:P3-C] This module is purely additive — no existing code is
deprecated or replaced. All fields are optional with sensible defaults.

Context flows through:
1. Tracing spans (injected as span attributes via get_context_attributes())
2. W3C tracestate (propagated via P3-B's set_tracestate_baggage())
3. Signal submissions (attached as metadata)

To find all P3-C code: grep -rn "CONTEXT_UPGRADE:P3-C" sdks/python/
"""
import contextvars
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Optional, Dict


@dataclass
class AgentBillContext:
    """Rich context for billing attribution and workflow correlation.
    
    Fields:
        scope_id: Billing scope for cost attribution (e.g., project, team, department)
        workflow_id: Business workflow grouping (e.g., "onboarding-flow")
        batch_id: Batch operation grouping for bulk cost aggregation
        correlation_id: Cross-system correlation identifier
        store_prompt: Whether to persist prompt content (default True for backward compat)
        metadata: Arbitrary key-value pairs propagated through calls
    """
    scope_id: Optional[str] = None
    workflow_id: Optional[str] = None
    batch_id: Optional[str] = None
    correlation_id: Optional[str] = None
    store_prompt: bool = True
    metadata: Dict[str, str] = field(default_factory=dict)


# Async-safe context storage
_context: contextvars.ContextVar[Optional[AgentBillContext]] = contextvars.ContextVar(
    'agentbill_context', default=None
)


def get_context() -> Optional[AgentBillContext]:
    """Get the current AgentBill context.
    
    Returns:
        Optional[AgentBillContext]: Current context or None
    """
    return _context.get()


def set_context(
    scope_id: Optional[str] = None,
    workflow_id: Optional[str] = None,
    batch_id: Optional[str] = None,
    correlation_id: Optional[str] = None,
    store_prompt: bool = True,
    metadata: Optional[Dict[str, str]] = None,
) -> AgentBillContext:
    """Set the AgentBill context explicitly.
    
    Args:
        scope_id: Billing scope for cost attribution
        workflow_id: Business workflow grouping
        batch_id: Batch operation grouping
        correlation_id: Cross-system correlation
        store_prompt: Whether to persist prompt content
        metadata: Arbitrary key-value pairs
        
    Returns:
        AgentBillContext: The newly set context
    """
    ctx = AgentBillContext(
        scope_id=scope_id,
        workflow_id=workflow_id,
        batch_id=batch_id,
        correlation_id=correlation_id,
        store_prompt=store_prompt,
        metadata=metadata or {},
    )
    _context.set(ctx)
    return ctx


def clear_context() -> None:
    """Clear the current AgentBill context."""
    _context.set(None)


@contextmanager
def context(
    scope_id: Optional[str] = None,
    workflow_id: Optional[str] = None,
    batch_id: Optional[str] = None,
    correlation_id: Optional[str] = None,
    store_prompt: bool = True,
    metadata: Optional[Dict[str, str]] = None,
):
    """Context manager for scoped AgentBill context.
    
    Restores the previous context when the block exits (even on exception).
    
    Example:
        >>> with context(scope_id="proj_abc", workflow_id="wf_123"):
        ...     result = wrapped_openai.chat(...)
    """
    previous = _context.get()
    ctx = AgentBillContext(
        scope_id=scope_id,
        workflow_id=workflow_id,
        batch_id=batch_id,
        correlation_id=correlation_id,
        store_prompt=store_prompt,
        metadata=metadata or {},
    )
    _context.set(ctx)
    try:
        yield ctx
    finally:
        _context.set(previous)


def get_context_attributes() -> Dict[str, str]:
    """Get the current context as span attributes.
    
    All attributes are prefixed with `agentbill.` namespace.
    Empty/None fields are omitted.
    Metadata keys are prefixed with `agentbill.meta.`.
    
    Returns:
        Dict[str, str]: Span attributes (empty dict if no context)
    """
    ctx = _context.get()
    if not ctx:
        return {}
    
    attrs: Dict[str, str] = {}
    if ctx.scope_id:
        attrs["agentbill.scope_id"] = ctx.scope_id
    if ctx.workflow_id:
        attrs["agentbill.workflow_id"] = ctx.workflow_id
    if ctx.batch_id:
        attrs["agentbill.batch_id"] = ctx.batch_id
    if ctx.correlation_id:
        attrs["agentbill.correlation_id"] = ctx.correlation_id
    if not ctx.store_prompt:
        attrs["agentbill.store_prompt"] = "false"
    for k, v in ctx.metadata.items():
        attrs[f"agentbill.meta.{k}"] = v
    
    return attrs


def get_context_baggage() -> Dict[str, str]:
    """Get the current context as tracestate baggage for W3C propagation.
    
    Only includes fields suitable for cross-service propagation:
    scope_id, workflow_id, correlation_id.
    
    batch_id, store_prompt, and metadata are local-only (not propagated).
    
    Returns:
        Dict[str, str]: Baggage key-value pairs (empty dict if no context)
    """
    ctx = _context.get()
    if not ctx:
        return {}
    
    baggage: Dict[str, str] = {}
    if ctx.scope_id:
        baggage["scope_id"] = ctx.scope_id
    if ctx.workflow_id:
        baggage["workflow_id"] = ctx.workflow_id
    if ctx.correlation_id:
        baggage["correlation_id"] = ctx.correlation_id
    
    return baggage
