"""Context variable for propagating user/client/presentation info into LLM calls."""
import contextvars
import uuid
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class LLMCallContext:
    user_id: Optional[uuid.UUID] = field(default=None)
    client_id: Optional[uuid.UUID] = field(default=None)
    presentation_id: Optional[uuid.UUID] = field(default=None)


_llm_context: contextvars.ContextVar[Optional[LLMCallContext]] = contextvars.ContextVar(
    "llm_call_context", default=None
)


def set_llm_context(ctx: LLMCallContext) -> None:
    _llm_context.set(ctx)


def get_llm_context() -> Optional[LLMCallContext]:
    return _llm_context.get()


def clear_llm_context() -> None:
    _llm_context.set(None)
