#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Langfuse 专题 · 步骤 5 · 修改后：给 langfuse_trace.py 补 Observation UUID，重启 worker。"""
from __future__ import annotations

import sys
import time

import _bootstrap  # noqa: F401
import _config as C
import _plugin_util as PU
from dify_api import Dify

LF_TRACE_PATH = (
    "/app/api/.venv/lib/python3.12/site-packages/dify_trace_langfuse/langfuse_trace.py"
)


def patch_container(d: Dify, container: str) -> None:
    content = PU.container_read(d, container, LF_TRACE_PATH)
    if PU.LF_TRACE_MARKER in content:
        print(container, "already patched")
        return
    new_content = PU.apply_langfuse_trace_patch(content)
    PU.container_write(d, container, LF_TRACE_PATH, new_content)
    out, err = d.sh(
        "docker exec %s python -m py_compile %s" % (container, LF_TRACE_PATH)
    )
    print(container, "py_compile:", out or err or "ok")


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        for ctr in (C.DIFY_WORKER, C.DIFY_API_CTR):
            patch_container(d, ctr)

        print("\n=== restart docker-worker-1 ===")
        out, err = d.sh("docker restart %s" % C.DIFY_WORKER, t=120)
        print(out or err)

        for i in range(15):
            out, _ = d.sh("docker logs --since 30s %s 2>&1 | tail -5" % C.DIFY_WORKER)
            if "ready" in (out or "").lower():
                print("worker ready")
                break
            time.sleep(2)
        else:
            print("[WARN] 未在日志中看到 celery ready，请手动 docker logs -f", C.DIFY_WORKER)

        for ctr in (C.DIFY_WORKER, C.DIFY_API_CTR):
            print(PU.grep_in_container(d, ctr, LF_TRACE_PATH, "auto-assign if missing"))
        return 0
    finally:
        d.close()


if __name__ == "__main__":
    sys.exit(main())