"""Module for request context."""

from __future__ import annotations

import contextvars
import uuid

request_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("request_id", default=None)
trace_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("trace_id", default=None)
span_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar("span_id", default=None)


def generate_request_id() -> str:
    return str(uuid.uuid4())[:8]


def set_request_context(
    *,
    request_id: str | None = None,
    trace_id: str | None = None,
    span_id: str | None = None,
) -> None:
    if request_id is not None:
        request_id_var.set(request_id)
    if trace_id is not None:
        trace_id_var.set(trace_id)
    if span_id is not None:
        span_id_var.set(span_id)


def clear_request_context() -> None:
    request_id_var.set(None)
    trace_id_var.set(None)
    span_id_var.set(None)


def get_request_id(*, fallback_generate: bool = True) -> str | None:
    value = request_id_var.get()
    if value:
        return value
    if not fallback_generate:
        return None
    value = generate_request_id()
    request_id_var.set(value)
    return value


def get_trace_ids() -> tuple[str | None, str | None]:
    return trace_id_var.get(), span_id_var.get()