#!/usr/bin/env python3
"""给 Dify 的 Langfuse 上报补上工具 Span。

现网的 dify-worker-1 已经打过。容器如果按镜像重建，这段会丢，再在 10.0.1.47 上执行本文件。
不重启 vLLM、New-API、dify-web-1。执行完会重启 dify-worker-1。
"""
from __future__ import annotations

import subprocess
from pathlib import Path

TRACE_PATH = "/app/api/.venv/lib/python3.12/site-packages/dify_trace_langfuse/langfuse_trace.py"
MARKER = "# AGENT_TOOL_SPANS_V2"
ANCHOR = """        self.add_generation(langfuse_generation_data)

    def moderation_trace(self, trace_info: ModerationTraceInfo):
"""
BLOCK = """        try:
            langfuse_generation_data.id = str(message_id)
        except Exception:
            pass

        self.add_generation(langfuse_generation_data)

        # AGENT_TOOL_SPANS_V2
        try:
            from sqlalchemy import select

            from core.db.session_factory import session_factory
            from models.model import MessageAgentThought

            with session_factory.create_session() as session:
                thoughts = session.scalars(
                    select(MessageAgentThought)
                    .where(MessageAgentThought.message_id == str(message_id))
                    .order_by(MessageAgentThought.position.asc())
                ).all()
                logger.info("agent tool spans count=%s message_id=%s", len(thoughts), message_id)
                for at in thoughts:
                    for tn in at.tools:
                        if not tn:
                            continue
                        tool_inp = at.tool_inputs_dict.get(tn)
                        outputs = at.tool_outputs_dict if isinstance(at.tool_outputs_dict, dict) else {}
                        tool_output = outputs.get(tn)
                        if tool_output in (None, "", {}):
                            tool_output = at.observation or None
                        meta = at.tool_meta.get(tn, {}) if isinstance(at.tool_meta, dict) else {}
                        time_cost_f = float(getattr(at, "latency", 0) or 0)
                        if not time_cost_f:
                            try:
                                time_cost_f = float((meta or {}).get("time_cost") or 0)
                            except (TypeError, ValueError):
                                time_cost_f = 0.0
                        if not time_cost_f:
                            time_cost_f = 0.2
                        start_t = at.created_at or trace_info.start_time
                        if start_t is not None and getattr(start_t, "tzinfo", None) is None:
                            start_t = start_t.replace(tzinfo=UTC)
                        end_t = start_t + timedelta(seconds=time_cost_f) if start_t else None
                        err = (meta or {}).get("error")
                        self.add_span(
                            LangfuseSpan(
                                name=tn,
                                input=tool_inp,
                                output=tool_output,
                                trace_id=trace_id,
                                parent_observation_id=str(message_id),
                                start_time=start_t,
                                end_time=end_t,
                                metadata={**metadata, "tool_name": tn, "time_cost": time_cost_f},
                                level=(LevelEnum.DEFAULT if not err else LevelEnum.ERROR),
                                status_message=str(err or ""),
                            )
                        )
        except Exception:
            logger.exception("agent tool spans patch failed")

    def moderation_trace(self, trace_info: ModerationTraceInfo):
"""
ID_OLD = 'id=data.get("id"),'
ID_NEW = 'id=data.get("id") or self._make_event_id(),'


def patch_text(text: str) -> str:
    if MARKER not in text:
        if ANCHOR not in text:
            raise SystemExit("找不到 message_trace 锚点，文件版本可能已变")
        text = text.replace(ANCHOR, BLOCK, 1)
    text = text.replace(ID_OLD, ID_NEW)
    return text


def write_container(container: str, text: str) -> None:
    remote = Path("/tmp/langfuse_trace.py")
    remote.write_text(text, encoding="utf-8")
    subprocess.check_call(["docker", "cp", str(remote), "%s:%s" % (container, TRACE_PATH)])
    check = subprocess.run(
        ["docker", "exec", container, "python", "-m", "py_compile", TRACE_PATH],
        capture_output=True,
        text=True,
    )
    if check.returncode != 0:
        raise SystemExit(check.stderr)
    print(container, "ok")


def main() -> None:
    raw = subprocess.check_output(["docker", "exec", "dify-worker-1", "cat", TRACE_PATH], text=True)
    patched = patch_text(raw)
    write_container("dify-worker-1", patched)
    write_container("dify-api-1", patched)
    subprocess.check_call(["docker", "restart", "dify-worker-1"])
    print("dify-worker-1 已重启")


if __name__ == "__main__":
    main()