"""APM span translator — OpenInference LLM spans -> Tencent APM format.

Mirrors frameworks-py/src/runner/apm_span_exporter.py.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
from opentelemetry.sdk.util.instrumentation import InstrumentationScope
from opentelemetry.trace import SpanKind
from opentelemetry.trace.status import Status, StatusCode

from .config import (
    APM_IDENTITY_LIBRARY_NAME,
    APM_IDENTITY_LIBRARY_VERSION,
    apm_identity_attrs,
)
from .llm_semconv import (
    OI_LLM_KIND_SET,
    OI_TO_GENAI_MIRROR,
    oi_kind_to_genai_span_kind,
    oi_kind_to_operation,
)


@dataclass
class TranslatorContext:
    is_root_span_id: Callable[[str], bool]
    root_span_has_children: Callable[[str], bool]


# LangGraph 控制流异常名（通过 interrupt() / Command 跳转触发），
# instrumentor 会经 on_chain_error 把对应 span 标成 ERROR，但这不是业务错误。
# 另外 GeneratorExit / CancelledError 是 interrupt 触发后流被拆除 / 取消时的
# teardown 信号（async generator aclose、task cancel），同样不是业务错误，
# 一并归一为非错误，避免 APM/langfuse 误报。
_INTERRUPT_EXCEPTION_NAMES = frozenset({
    "GraphInterrupt",
    "NodeInterrupt",
    "ParentCommand",
    "GeneratorExit",
    "CancelledError",
})

# CrewAI 的人工反馈暂停信号：crewai.flow.async_feedback.types.HumanFeedbackPending
# 是 ``Exception`` 的子类（源码注释 noqa N818 — Not an error, a control flow signal），
# 用于 @human_feedback flow 暂停执行、持久化状态并把异常对象返回给调用者，并非业务错误。
# openinference CrewAIInstrumentor 仍会把它经 record_exception 标到 span 上导致误报。
# 与 LangGraph interrupt 不同，这类名字本身虽 CrewAI 专属，但为保守起见仅在 span 来自
# crewai instrumentation scope 时才归一，避免误伤其它来源的同名异常。
_CREWAI_INTERRUPT_EXCEPTION_NAMES = frozenset({
    "HumanFeedbackPending",
})


def _attr(span: ReadableSpan, key: str) -> str | None:
    v = span.attributes.get(key) if span.attributes else None
    if v is None:
        return None
    return str(v)


def _iter_exception_type_names(span: ReadableSpan):
    """产出 span 上记录的异常类名（去掉模块前缀）。

    覆盖两种惯例：异常类型写在 ``exception.type`` / ``error.type`` 属性上，
    或作为 OTel 标准 "exception" event 的 ``exception.type`` 记录。
    """
    attrs = span.attributes or {}
    for key in ("exception.type", "error.type"):
        v = attrs.get(key)
        if isinstance(v, str) and v:
            yield v.rsplit(".", 1)[-1]
    for ev in getattr(span, "events", None) or ():
        if getattr(ev, "name", None) != "exception":
            continue
        ev_attrs = getattr(ev, "attributes", None) or {}
        v = ev_attrs.get("exception.type")
        if isinstance(v, str) and v:
            yield v.rsplit(".", 1)[-1]


def _is_crewai_scope(span: ReadableSpan) -> bool:
    """span 是否来自 crewai instrumentation scope。

    覆盖 openinference 的 ``openinference.instrumentation.crewai`` 与 crewai 原生埋点，
    统一按 scope 名是否包含 "crewai"（忽略大小写）判断。
    """
    scope = getattr(span, "instrumentation_scope", None)
    name = getattr(scope, "name", None)
    return isinstance(name, str) and "crewai" in name.lower()


def _matches_interrupt_name(name: str, crewai_scope: bool) -> bool:
    """判断异常类名（已去掉模块前缀）是否为控制流暂停信号。

    LangGraph 控制流异常全局归一；CrewAI HumanFeedbackPending 仅在 crewai scope 内归一。
    """
    if name in _INTERRUPT_EXCEPTION_NAMES:
        return True
    if crewai_scope and name in _CREWAI_INTERRUPT_EXCEPTION_NAMES:
        return True
    return False


def _is_interrupt_span(span: ReadableSpan) -> bool:
    """检测 span 是否由控制流暂停信号触发（非业务错误）。

    依据 OpenTelemetry / OpenInference 把异常类名写到 ``exception.type`` 或
    ``error.type`` 这一惯例。匹配时支持 ``module.GraphInterrupt`` 这类全限定名。
    """
    crewai_scope = _is_crewai_scope(span)
    return any(
        _matches_interrupt_name(n, crewai_scope)
        for n in _iter_exception_type_names(span)
    )


def _filter_interrupt_exception_events(span: ReadableSpan):
    """剔除 span 上记录的控制流暂停异常 event，返回过滤后的 event 元组。

    APM/langfuse 平台会依据 span 上 ``record_exception()`` 产生的 ``exception``
    event 判定 ``error=true`` / ``statusCode=ERROR`` 并展示堆栈，仅把 OTel status
    归一为 UNSET 不足以消除误报。因此对 interrupt span 同步剔除对应的 exception
    event。仅剔除类型命中 interrupt 名单的 event，保留其它真实异常 event。

    返回 None 表示没有任何 event 被剔除（调用方应沿用原始 events）。
    """
    crewai_scope = _is_crewai_scope(span)
    kept = []
    dropped = False
    for ev in getattr(span, "events", None) or ():
        if getattr(ev, "name", None) == "exception":
            ev_attrs = getattr(ev, "attributes", None) or {}
            etype = ev_attrs.get("exception.type")
            if (
                isinstance(etype, str)
                and etype
                and _matches_interrupt_name(etype.rsplit(".", 1)[-1], crewai_scope)
            ):
                dropped = True
                continue
        kept.append(ev)
    return tuple(kept) if dropped else None


def _is_oi_llm_span(span: ReadableSpan) -> bool:
    kind = _attr(span, "openinference.span.kind")
    return kind is not None and kind in OI_LLM_KIND_SET


def _is_trace_entry(span: ReadableSpan, ctx: TranslatorContext) -> bool:
    parent = span.parent
    if parent is None:
        return True
    try:
        pid_int = parent.span_id
    except AttributeError:
        return False
    pid_hex = format(pid_int, "016x")
    return ctx.is_root_span_id(pid_hex)


def _genai_mirror(src: dict[str, Any]) -> dict[str, Any]:
    out: dict[str, Any] = {}
    for dest_key, src_key, coerce in OI_TO_GENAI_MIRROR:
        if dest_key in src:
            continue
        v = src.get(src_key)
        if v is None:
            continue
        if coerce == "number":
            try:
                n = float(v)
            except (TypeError, ValueError):
                continue
            if n != n or n in (float("inf"), float("-inf")):
                continue
            out[dest_key] = int(n) if n.is_integer() else n
        else:
            out[dest_key] = str(v)

    oi_kind = src.get("openinference.span.kind")
    if isinstance(oi_kind, str):
        mapped = oi_kind_to_operation(oi_kind)
        if "llm.request.type" not in src:
            out["llm.request.type"] = mapped
        if "gen_ai.operation.name" not in src:
            out["gen_ai.operation.name"] = mapped
        if "gen_ai.span.kind" not in src:
            out["gen_ai.span.kind"] = oi_kind_to_genai_span_kind(oi_kind)

    return out


class _TranslatedSpan:
    __slots__ = ("_src", "_attributes", "_status", "_scope", "_events")

    def __init__(
        self,
        src: ReadableSpan,
        attributes: dict[str, Any],
        status: Status,
        scope: InstrumentationScope,
        events: Any = None,
    ) -> None:
        self._src = src
        self._attributes = attributes
        self._status = status
        self._scope = scope
        self._events = events

    @property
    def name(self) -> str:
        return self._src.name

    @property
    def kind(self) -> SpanKind:
        return SpanKind.INTERNAL

    @property
    def context(self):
        return self._src.context

    def get_span_context(self):
        return self._src.get_span_context()

    @property
    def parent(self):
        return self._src.parent

    @property
    def start_time(self) -> int | None:
        return self._src.start_time

    @property
    def end_time(self) -> int | None:
        return self._src.end_time

    @property
    def status(self) -> Status:
        return self._status

    @property
    def attributes(self) -> dict[str, Any]:
        return self._attributes

    @property
    def instrumentation_scope(self) -> InstrumentationScope:
        return self._scope

    @property
    def instrumentation_info(self) -> InstrumentationScope:
        return self._scope

    @property
    def events(self):
        return self._src.events if self._events is None else self._events

    @property
    def links(self):
        return self._src.links

    @property
    def resource(self):
        return self._src.resource

    @property
    def dropped_attributes(self) -> int:
        return getattr(self._src, "dropped_attributes", 0)

    @property
    def dropped_events(self) -> int:
        return getattr(self._src, "dropped_events", 0)

    @property
    def dropped_links(self) -> int:
        return getattr(self._src, "dropped_links", 0)

    def to_json(self, *args: Any, **kwargs: Any) -> str:
        return self._src.to_json(*args, **kwargs)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._src, name)


def _translate(span: ReadableSpan, ctx: TranslatorContext) -> ReadableSpan:
    src_attrs: dict[str, Any] = dict(span.attributes or {})
    model = src_attrs.get("llm.model_name") or "unknown"
    provider = src_attrs.get("llm.provider") or "unknown"
    oi_kind = src_attrs.get("openinference.span.kind") or "unknown"

    # APM 平台对普通 attribute 有 1 KB 的截断限制，但以 "langfuse." 开头的
    # attribute 不受限制。所有 span attribute 统一加 "langfuse." 前缀以避免截断。
    merged: dict[str, Any] = {f"langfuse.{k}": v for k, v in src_attrs.items()}
    merged.setdefault("custom_key_1", str(model))
    merged.setdefault("custom_key_2", str(provider))
    merged.setdefault("custom_key_3", str(oi_kind))

    merged.update(_genai_mirror(src_attrs))

    if "gen_ai.is_entry" not in merged and _is_trace_entry(span, ctx):
        merged["gen_ai.is_entry"] = True

    for k, v in apm_identity_attrs().items():
        merged.setdefault(k, v)

    status = span.status or Status(StatusCode.UNSET)
    if status.status_code == StatusCode.OK:
        status = Status(StatusCode.UNSET)
    events_override = None
    # LangGraph interrupt 不是真错误，把 ERROR 归一回 UNSET，并打标方便筛选。
    if status.status_code == StatusCode.ERROR and _is_interrupt_span(span):
        status = Status(StatusCode.UNSET)
        merged["agent.interrupt"] = True
        merged["langfuse.agent.interrupt"] = True
        # 同步剔除 interrupt 异常 event，避免 APM 仍据此判 error 并展示堆栈。
        events_override = _filter_interrupt_exception_events(span)

    scope = InstrumentationScope(
        name=APM_IDENTITY_LIBRARY_NAME,
        version=APM_IDENTITY_LIBRARY_VERSION,
    )

    return _TranslatedSpan(
        src=span,
        attributes=merged,
        status=status,
        scope=scope,
        events=events_override,
    )


def _prefix_only(span: ReadableSpan, ctx: TranslatorContext) -> ReadableSpan:
    """对非 OI LLM span 只做 attribute 前缀处理，不注入 APM 专用字段。
    例外：root span 若有子 span（说明该请求有实际的追踪活动），标记 gen_ai.is_entry。
    """
    src_attrs: dict[str, Any] = dict(span.attributes or {})
    prefixed: dict[str, Any] = {f"langfuse.{k}": v for k, v in src_attrs.items()}
    span_id = format(span.get_span_context().span_id, "016x")
    if ctx.is_root_span_id(span_id) and ctx.root_span_has_children(span_id):
        prefixed["gen_ai.is_entry"] = True

    status = span.status or Status(StatusCode.UNSET)
    events_override = None
    if status.status_code == StatusCode.ERROR and _is_interrupt_span(span):
        status = Status(StatusCode.UNSET)
        prefixed["agent.interrupt"] = True
        prefixed["langfuse.agent.interrupt"] = True
        # 同步剔除 interrupt 异常 event，避免 APM 仍据此判 error 并展示堆栈。
        events_override = _filter_interrupt_exception_events(span)

    return _TranslatedSpan(
        src=span,
        attributes=prefixed,
        status=status,
        scope=span.instrumentation_scope,
        events=events_override,
    )


class ApmSpanExporter(SpanExporter):
    def __init__(self, inner: SpanExporter, ctx: TranslatorContext) -> None:
        self._inner = inner
        self._ctx = ctx

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        translated = [
            _translate(s, self._ctx) if _is_oi_llm_span(s) else _prefix_only(s, self._ctx)
            for s in spans
        ]
        return self._inner.export(translated)

    def shutdown(self) -> None:
        self._inner.shutdown()

    def force_flush(self, timeout_millis: int = 30_000) -> bool:
        try:
            return self._inner.force_flush(timeout_millis)
        except AttributeError:
            return True