"""Context-scoped feature metadata using ContextVar."""

from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar

_feature_metadata: ContextVar[dict[str, str]] = ContextVar("caliper_feature_metadata", default={})
_linked_request_id: ContextVar[str | None] = ContextVar("caliper_linked_request_id", default=None)


@contextmanager
def features(*, previous_request: str | None = None, **kwargs: str) -> Generator[None, None, None]:
    """Attach metadata to all LLM calls within this block. Nestable — inner merges with outer.

    Use previous_request to link all calls in this block to a prior request by its caliper request ID.
    """
    previous_meta = _feature_metadata.get()
    merged = {**previous_meta, **{k: str(v) for k, v in kwargs.items()}}
    meta_token = _feature_metadata.set(merged)

    # Only set linked_request_id if explicitly provided; otherwise inherit from outer scope
    link_token = None
    if previous_request is not None:
        link_token = _linked_request_id.set(previous_request)

    try:
        yield
    finally:
        _feature_metadata.reset(meta_token)
        if link_token is not None:
            _linked_request_id.reset(link_token)


def get_feature_metadata() -> dict[str, str]:
    """Return a copy of the current feature metadata."""
    return dict(_feature_metadata.get())


def get_linked_request_id() -> str | None:
    """Return the linked request ID from the current features context, or None."""
    return _linked_request_id.get()
