#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""请求/工作流上下文：基于 ``contextvars`` 的双层 ID 传播。

设计上把"一次用户请求"和"一次重型工作流"区分开：

- ``request_id``：会话域。所有进 API 的请求都要带，由 middleware 在入口生成或
  从 ``X-Request-Id`` 头继承。chat 这条同步链路靠它走完一生。
- ``trace_id``：工作流域。只有 ``ConsciousnessNode`` 决定启动重型任务时才生成，
  挂到 ``KiloStarWorkflow`` 上。trace_id 应能追溯回触发它的 request_id（前者
  通过显式参数传入，后者从 contextvars 读取）。

为什么用 ``contextvars`` 而不是参数透传：

1. ``contextvars`` 在 ``asyncio`` 协程间天然继承，不会跨协程串味；
2. ``loguru`` 的 ``patcher`` 钩子可以把它变成日志切面，业务代码不需要在每条
   ``logger.info`` 上手动 ``.bind(trace_id=...)``；
3. 跨 asyncio.create_task 边界时也会自动继承当前 context，符合单进程 actor 模型
   的语义 —— 一次任务链路的 ID 与其入口请求一致，不必显式透传。
"""

from __future__ import annotations

import uuid
from contextlib import contextmanager
from contextvars import ContextVar, Token
from typing import Iterator


_request_id_var: ContextVar[str] = ContextVar("kilostar_request_id", default="")
_trace_id_var: ContextVar[str] = ContextVar("kilostar_trace_id", default="")


def get_request_id() -> str:
    """返回当前协程的 ``request_id``，未绑定时返回空串。"""
    return _request_id_var.get()


def get_trace_id() -> str:
    """返回当前协程的 ``trace_id``，未绑定时返回空串。"""
    return _trace_id_var.get()


def bind_request_id(request_id: str) -> Token:
    """直接绑定 ``request_id`` 到当前 context，返回 token 以便 ``reset`` 还原。

    返回的 ``Token`` 只能在与 ``set`` 同一线程/协程中传给 ``reset``，否则会抛
    ``ValueError``。一般情况下推荐用 ``request_id_scope`` 上下文管理器代替。
    """
    return _request_id_var.set(request_id)


def bind_trace_id(trace_id: str) -> Token:
    """直接绑定 ``trace_id`` 到当前 context，返回 token 以便 ``reset`` 还原。"""
    return _trace_id_var.set(trace_id)


def reset_request_id(token: Token) -> None:
    _request_id_var.reset(token)


def reset_trace_id(token: Token) -> None:
    _trace_id_var.reset(token)


@contextmanager
def request_id_scope(request_id: str) -> Iterator[str]:
    """``with`` 范围内绑定 request_id，退出自动还原。"""
    token = _request_id_var.set(request_id)
    try:
        yield request_id
    finally:
        _request_id_var.reset(token)


@contextmanager
def trace_id_scope(trace_id: str) -> Iterator[str]:
    """``with`` 范围内绑定 trace_id，退出自动还原。"""
    token = _trace_id_var.set(trace_id)
    try:
        yield trace_id
    finally:
        _trace_id_var.reset(token)


def new_request_id(prefix: str = "req") -> str:
    """生成一个新的 request_id：``<prefix>-<uuid4 hex>``。"""
    return f"{prefix}-{uuid.uuid4().hex}"


def snapshot() -> dict[str, str]:
    """返回当前上下文 ID 的快照，便于日志/追踪或结构化写库。"""
    return {
        "request_id": _request_id_var.get(),
        "trace_id": _trace_id_var.get(),
    }