"""Decorator utilities for agentbudget."""

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

from agentbudget.session import CostSession

F = TypeVar("F", bound=Callable[..., Any])


def track(
    label: str = "",
    budget: Optional[float] = None,
    strict: bool = False,
    verbose: bool = True,
) -> Callable[[F], F]:
    """Decorator that wraps a function in a CostSession for automatic cost tracking.

    The wrapped function must return an LLM API response object. The decorator
    creates a CostSession for each call and tracks the returned response.

    Args:
        label: Step label used in the session summary. Defaults to the function name.
        budget: Optional spending limit in USD per call.
        strict: If True, raise BudgetExceededError when budget is exceeded.
        verbose: If True (default), print the session summary after each call.

    Example:
        @track(label="generate-ideas", budget=0.02)
        def generate_ideas(topic: str):
            return openai_client.chat.completions.create(...)
    """
    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            step_label = label or func.__name__
            session_name = func.__name__
            with CostSession(name=session_name, budget=budget, strict=strict, verbose=verbose) as session:
                response = func(*args, **kwargs)
                if response is None:
                    raise ValueError(
                        f"[agentbudget] @track: function '{func.__name__}' returned None. "
                        "The decorated function must return an LLM API response object."
                    )
                session.track(response, label=step_label)
            return response
        return wrapper  # type: ignore[return-value]
    return decorator
