"""Context management for request-scoped logging fields."""

import threading
from contextlib import contextmanager
from typing import Any


_context = threading.local()


def get_context() -> dict[str, Any]:
    """Get the current thread-local context dict."""
    if not hasattr(_context, "fields"):
        _context.fields = {}
    return _context.fields


def update(**kwargs) -> dict[str, Any]:
    """Add fields to the current context. Returns the updated context."""
    ctx = get_context()
    ctx.update(kwargs)
    return ctx


def clear() -> None:
    """Clear all context fields for the current thread."""
    _context.fields = {}


@contextmanager
def context(**kwargs):
    """Context manager that adds fields to logs within the block.

    >>> with context(request_id="r1", user_id="u42"):
    ...     log.info("Processing")  # includes request_id + user_id
    """
    old = dict(get_context())
    get_context().update(kwargs)
    try:
        yield get_context()
    finally:
        _context.fields = old
