"""Agent context propagator — Python parallel of Node's AgentContextPropagator.

为每个 span 在 onStart 阶段统一注入：
- ``agent.run_id`` = 该 span 所属 trace 的 ``traceId``（与 traceId 对齐，
  上报到 APM / langfuse 后体现为 ``langfuse.agent.run_id == traceId``）。
- ``agent.conversation_id`` = adapter 通过 ``set_agent_conversation_id`` 设置的会话 ID。

注意：业务层 ``ctx.run_id`` 仍由 runtime.py 用 ``uuid4()`` 生成，不受影响；
本 propagator 只影响上报到 OTel exporter 的 attribute。
"""
from __future__ import annotations

from contextvars import ContextVar
from typing import Optional

from opentelemetry.context import Context
from opentelemetry.sdk.trace import ReadableSpan, Span
from opentelemetry.sdk.trace.export import SpanProcessor


# 由 adapter 在 attach 完 root span context 之后立刻 set，handler 完成后 reset。
# 用 ContextVar 而不是模块全局，自然适配 asyncio 的 Task 隔离。
_conversation_id_var: ContextVar[Optional[str]] = ContextVar(
    "_eo_agent_conversation_id", default=None
)


def set_agent_conversation_id(conversation_id: Optional[str]):
    """Adapter 调用：设置当前请求的 conversation_id，返回 token 供 reset 用。

    Returns:
        ContextVar token; 调用方应在请求结束时执行
        ``_conversation_id_var.reset(token)``。也可直接调用
        :func:`reset_agent_conversation_id` 传入该 token。
    """
    return _conversation_id_var.set(conversation_id)


def reset_agent_conversation_id(token) -> None:
    try:
        _conversation_id_var.reset(token)
    except Exception:
        # token 失效或跨 context 复位等异常一律忽略，best-effort 清理。
        pass


class AgentContextPropagator(SpanProcessor):
    """SpanProcessor that injects ``agent.run_id`` (= traceId) and
    ``agent.conversation_id`` into every span on start.
    """

    def on_start(
        self, span: Span, parent_context: Optional[Context] = None
    ) -> None:  # noqa: D401
        try:
            ctx = span.get_span_context()
            # 16-byte trace id → 32 char hex，对齐 OTel/langfuse 显示格式。
            trace_id_hex = format(ctx.trace_id, "032x")
            if trace_id_hex and trace_id_hex != "0" * 32:
                span.set_attribute("agent.run_id", trace_id_hex)
        except Exception:
            pass

        conv_id = _conversation_id_var.get()
        if conv_id:
            try:
                span.set_attribute("agent.conversation_id", conv_id)
            except Exception:
                pass

    def on_end(self, span: ReadableSpan) -> None:  # noqa: D401
        return None

    def shutdown(self) -> None:  # noqa: D401
        return None

    def force_flush(self, timeout_millis: int = 30000) -> bool:  # noqa: D401
        return True