#!/usr/bin/env python3
"""注册 getServerStatus，并挂到「运维巡检 Agent」。

现网已经执行过。再跑时，同名工具和同名 Agent 已存在就只打印编号，不删除、不重建。
在 10.0.1.47 上执行。Windows 上传后先去掉 CR：sed -i 's/\r$//' 本文件。
"""
from __future__ import annotations

import base64
import json
from http.cookiejar import CookieJar
from pathlib import Path
from urllib import request

BASE = "http://127.0.0.1:8088/console/api"
PROVIDER = "server_status"
TOOL_NAME = "getServerStatus"
APP_NAME = "运维巡检 Agent"
AGENT_ID = "01a124c7-9fce-70a8-9efa-41b2a6f5b7e7"
APP_ID = "e4d614f6-de22-4f4c-b58e-caeb0504b121"
PROVIDER_ID = "a69388a5-c02f-4bf1-8c28-1f20ee0687e6"
SCHEMA_PATH = Path(__file__).resolve().parents[1] / "config" / "server-status-openapi.yaml"

SYSTEM_PROMPT = """你是企业内部运维助手。当用户询问某台服务器的状态、CPU、内存、GPU 时，
必须调用工具 getServerStatus（参数 server_name），并根据工具返回结果用中文格式化汇报。

汇报格式示例：
- 服务器：{server}
- CPU：{cpu}
- 内存：{memory}
- GPU 利用率：{gpu_util}
- GPU 显存：{gpu_memory}

不要编造指标；工具失败时如实说明。"""

ADD_BODY = {
    "provider": PROVIDER,
    "credentials": {"auth_type": "none"},
    "schema_type": "openapi",
    "schema": SCHEMA_PATH.read_text(encoding="utf-8"),
    "icon": {"background": "#E8F5E9", "content": "S"},
    "privacy_policy": "",
    "custom_disclaimer": "",
    "labels": [],
}

CREATE_AGENT_BODY = {
    "name": APP_NAME,
    "description": "带 HTTP 工具的运维巡检 Agent，查询服务器状态",
    "icon_type": "emoji",
    "icon": "🖥️",
    "icon_background": "#E8F5E9",
    "role": "ops",
}

TOOL_ENTRY = {
    "enabled": True,
    "provider_type": "api",
    "provider_id": PROVIDER_ID,
    "tool_name": TOOL_NAME,
    "credential_type": "unauthorized",
    "credential_ref": None,
    "runtime_parameters": {},
}


class Dify:
    def __init__(self) -> None:
        self.jar = CookieJar()
        self.opener = request.build_opener(request.HTTPCookieProcessor(self.jar))
        self.csrf = ""

    def login(self) -> None:
        self.call(
            "POST",
            "/login",
            {
                "email": "admin@example.com",
                "password": base64.b64encode(b"Dify-rag-2026").decode(),
                "language": "zh-Hans",
                "remember_me": True,
            },
        )
        for cookie in self.jar:
            if cookie.name == "csrf_token":
                self.csrf = cookie.value
        if not self.csrf:
            raise SystemExit("没有拿到 csrf_token")

    def call(self, method: str, path: str, payload: dict | None = None) -> tuple[int, dict | list]:
        data = None if payload is None else json.dumps(payload).encode()
        req = request.Request(BASE + path, data=data, method=method)
        if payload is not None:
            req.add_header("Content-Type", "application/json")
        if self.csrf:
            req.add_header("X-CSRF-Token", self.csrf)
        with self.opener.open(req, timeout=60) as resp:
            raw = resp.read().decode()
            body = json.loads(raw) if raw else {}
            return resp.status, body


def enable_trace(dify: Dify, app_id: str) -> None:
    env: dict[str, str] = {}
    for line in open("/usr/cluster/observability/langfuse/.env", encoding="utf-8"):
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            key, value = line.split("=", 1)
            env[key] = value.strip().strip('"')
    dify.call(
        "POST",
        "/apps/%s/trace-config" % app_id,
        {
            "tracing_provider": "langfuse",
            "tracing_config": {
                "public_key": env["LANGFUSE_INIT_PROJECT_PUBLIC_KEY"],
                "secret_key": env["LANGFUSE_INIT_PROJECT_SECRET_KEY"],
                "host": "http://172.17.0.1:3100",
            },
        },
    )
    dify.call(
        "POST",
        "/apps/%s/trace" % app_id,
        {"enabled": True, "tracing_provider": "langfuse"},
    )
    print("追踪已打开", app_id)


def main() -> None:
    dify = Dify()
    dify.login()
    _status, tools = dify.call("GET", "/workspaces/current/tools/api")
    names = [(item.get("name") or item.get("provider")) for item in tools]
    if PROVIDER in names:
        print("工具已在", PROVIDER, PROVIDER_ID)
    else:
        status, _created = dify.call("POST", "/workspaces/current/tool-provider/api/add", ADD_BODY)
        print("已注册", PROVIDER, status)

    _status, listed = dify.call("GET", "/agent?page=1&limit=50")
    found = None
    for item in listed.get("data") or []:
        if item.get("name") == APP_NAME:
            found = item
            break
    if found:
        print("Agent 已在", found.get("agent_id") or found.get("id"), "app", APP_ID)
        enable_trace(dify, APP_ID)
        return

    status, created = dify.call("POST", "/agent", CREATE_AGENT_BODY)
    agent_id = created.get("agent_id") or created.get("id")
    print("已创建", status, agent_id)
    _status, state = dify.call("GET", "/agent/%s/composer" % agent_id)
    soul = state["agent_soul"]
    soul.setdefault("prompt", {})["system_prompt"] = SYSTEM_PROMPT
    soul.setdefault("tools", {})["dify_tools"] = [TOOL_ENTRY]
    soul.setdefault("tools", {})["cli_tools"] = soul["tools"].get("cli_tools") or []
    model = soul.get("model") if isinstance(soul.get("model"), dict) else {}
    model.update(
        {
            "credential_ref": None,
            "model": "qwen3.6-27b",
            "model_provider": "langgenius/openai_api_compatible/openai_api_compatible",
            "plugin_id": "langgenius/openai_api_compatible",
        }
    )
    soul["model"] = model
    dify.call(
        "PUT",
        "/agent/%s/composer" % agent_id,
        {
            "variant": "agent_app",
            "save_strategy": "save_to_current_version",
            "agent_soul": soul,
            "version_note": "挂上 getServerStatus",
        },
    )
    dify.call("POST", "/agent/%s/publish" % agent_id, {"version_note": "挂上 getServerStatus"})
    print("已发布", agent_id)
    enable_trace(dify, state.get("app_id") or state.get("backing_app_id") or APP_ID)


if __name__ == "__main__":
    main()