"""
MuBit Auto-Capture Decorator.

Function decorator to group LLM calls into logical spans.
"""

import functools
import inspect
from typing import Any, Callable, Optional

from mubit.auto._context import Span, get_current_span, set_current_span


def observe(
    name: Optional[str] = None,
    session_id: Optional[str] = None,
    agent_id: Optional[str] = None,
    user_id: Optional[str] = None,
) -> Callable:
    """
    Decorator to trace a function and group enclosed LLM calls.
    
    Args:
        name: Name of the span (defaults to function name).
        session_id: Override session/run ID for this span.
        agent_id: Override agent ID for this span.
        user_id: Override user ID.
    """
    def decorator(fn: Callable) -> Callable:
        obs_name = name or fn.__name__

        @functools.wraps(fn)
        def sync_wrapper(*args, **kwargs):
            parent = get_current_span()
            
            # Inherit IDs from parent if not specified
            sid = session_id or (parent.session_id if parent else "")
            aid = agent_id or (parent.agent_id if parent else "auto")
            uid = user_id or (parent.user_id if parent else "")
            
            span = Span(
                name=obs_name,
                parent_id=parent.trace_id if parent else None,
                session_id=sid,
                agent_id=aid,
                user_id=uid,
            )
            
            token = set_current_span(span)
            try:
                result = fn(*args, **kwargs)
                # Store result on span? Only if needed for outcomes later.
                # For now, just return unmodified.
                return result
            finally:
                set_current_span(parent)  # restore via set or reset?
                # Using set allows restoring None correctly.
                # reset(token) is cleaner but contextvars.Token is what we have.
                # Let's use contextvar.reset(token) as intended.
                # But wait, set_current_span returns token.
                # Correct usage: _current_span.reset(token)
                # We need access to the var itself.
                # Let's import it or expose a reset helper.
                from mubit.auto._context import _current_span
                _current_span.reset(token)

        @functools.wraps(fn)
        async def async_wrapper(*args, **kwargs):
            parent = get_current_span()
            
            sid = session_id or (parent.session_id if parent else "")
            aid = agent_id or (parent.agent_id if parent else "auto")
            uid = user_id or (parent.user_id if parent else "")
            
            span = Span(
                name=obs_name,
                parent_id=parent.trace_id if parent else None,
                session_id=sid,
                agent_id=aid,
                user_id=uid,
            )
            
            token = set_current_span(span)
            try:
                result = await fn(*args, **kwargs)
                return result
            finally:
                from mubit.auto._context import _current_span
                _current_span.reset(token)

        if inspect.iscoroutinefunction(fn):
            return async_wrapper
        return sync_wrapper

    return decorator
