"""Request-scoped context propagated into logs and traces (架构文档 §15)."""

from __future__ import annotations

import uuid
from contextvars import ContextVar, Token
from dataclasses import dataclass

_request_id: ContextVar[str | None] = ContextVar("meida_request_id", default=None)
_trace_id: ContextVar[str | None] = ContextVar("meida_trace_id", default=None)
_user_id: ContextVar[int | None] = ContextVar("meida_user_id", default=None)
_task_id: ContextVar[int | None] = ContextVar("meida_task_id", default=None)


@dataclass(slots=True)
class ContextTokens:
    request_id: Token[str | None]
    trace_id: Token[str | None]
    user_id: Token[int | None]
    task_id: Token[int | None]


def new_request_id() -> str:
    return uuid.uuid4().hex


def bind_context(
    *,
    request_id: str | None = None,
    trace_id: str | None = None,
    user_id: int | None = None,
    task_id: int | None = None,
) -> ContextTokens:
    return ContextTokens(
        request_id=_request_id.set(request_id if request_id is not None else _request_id.get()),
        trace_id=_trace_id.set(trace_id if trace_id is not None else _trace_id.get()),
        user_id=_user_id.set(user_id if user_id is not None else _user_id.get()),
        task_id=_task_id.set(task_id if task_id is not None else _task_id.get()),
    )


def reset_context(tokens: ContextTokens) -> None:
    _request_id.reset(tokens.request_id)
    _trace_id.reset(tokens.trace_id)
    _user_id.reset(tokens.user_id)
    _task_id.reset(tokens.task_id)


def get_request_id() -> str | None:
    return _request_id.get()


def get_trace_id() -> str | None:
    return _trace_id.get()


def get_user_id() -> int | None:
    return _user_id.get()


def get_task_id() -> int | None:
    return _task_id.get()


def context_snapshot() -> dict[str, str | int]:
    snapshot: dict[str, str | int] = {}
    for key, value in (
        ("request_id", _request_id.get()),
        ("trace_id", _trace_id.get()),
        ("user_id", _user_id.get()),
        ("task_id", _task_id.get()),
    ):
        if value is not None:
            snapshot[key] = value
    return snapshot