from contextvars import ContextVar
from dataclasses import dataclass, replace
from contextlib import contextmanager
from typing import Optional


@dataclass(frozen=True)
class LLMCallContext:
    user_id: Optional[str] = None
    focus_group_id: Optional[str] = None
    persona_id: Optional[str] = None
    feature: str = "other"
    task_id: Optional[str] = None


_ctx: ContextVar[LLMCallContext] = ContextVar("llm_call_context", default=LLMCallContext())


def current_context() -> LLMCallContext:
    return _ctx.get()


def set_llm_context(**overrides) -> None:
    """Mutate the LLM context for the current asyncio task without cleanup.

    Use this at service entry points where the feature/focus_group_id/persona_id
    should persist for the duration of the whole async call tree (including sub-awaits).
    The change lives until the asyncio Task ends or is overridden again.

    Unlike llm_context(), this does NOT restore the previous value on exit — suitable
    for top-level service calls, not for re-entrant helpers.
    """
    prev = _ctx.get()
    _ctx.set(replace(prev, **overrides))


@contextmanager
def llm_context(**overrides):
    """Context manager that sets LLM call attribution metadata.

    Usage:
        with llm_context(user_id="abc", focus_group_id="xyz", feature="moderator"):
            await LLMService.generate_content(...)

    Overrides stack — inner contexts extend (not replace) outer ones.
    Safe across asyncio tasks and run_coroutine_threadsafe hops because
    ContextVar inherits context on task creation / thread submission.
    """
    prev = _ctx.get()
    token = _ctx.set(replace(prev, **overrides))
    try:
        yield
    finally:
        _ctx.reset(token)