"""
MuBit Auto-Capture Context Management.

This module provides the contextvars-based span tree that links LLM calls to their
enclosing logical step (function decorated with @observe) and manages capture state.
"""

import contextvars
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

# Current active span (if any)
_current_span: contextvars.ContextVar[Optional["Span"]] = contextvars.ContextVar(
    "mubit_auto_span", default=None
)

# Global capture toggle
_capture_enabled: contextvars.ContextVar[bool] = contextvars.ContextVar(
    "mubit_auto_capture", default=True
)


@dataclass
class Span:
    """A logical unit of work (trace span) that groups LLM interactions."""

    trace_id: str = field(default_factory=lambda: uuid.uuid4().hex[:16])
    parent_id: Optional[str] = None
    name: str = ""
    agent_id: str = "auto"
    session_id: str = ""
    user_id: str = ""
    start_time: float = field(default_factory=time.monotonic)
    items: List[Dict[str, Any]] = field(default_factory=list)

    def add_item(self, item: Dict[str, Any]) -> None:
        """Link an ingest item to this span."""
        # Note: metadata_json manipulation happens at item build time,
        # but we track items here for future aggregation if needed.
        self.items.append(item)


def get_current_span() -> Optional[Span]:
    """Get the currently active span from contextvars."""
    return _current_span.get()


def set_current_span(span: Optional[Span]) -> contextvars.Token:
    """Set the currently active span."""
    return _current_span.set(span)


def is_capture_enabled() -> bool:
    """Check if auto-capture is currently enabled."""
    return _capture_enabled.get()


@contextmanager
def no_capture():
    """Context manager to temporarily disable auto-capture."""
    token = _capture_enabled.set(False)
    try:
        yield
    finally:
        _capture_enabled.reset(token)
