#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业 1 · 步骤 2：在 Dify 工作区注册「查询服务器状态」自定义 API 工具。"""
from __future__ import annotations

import sys
from pathlib import Path

import _bootstrap  # noqa: F401
import _config as C
from dify_api import Dify, pp

OPENAPI = Path(__file__).resolve().parents[2] / "patches_cursor" / "作业1_带工具调用的Agent_服务器状态" / "作业1_server_status_openapi.yaml"


def load_schema() -> str:
    return OPENAPI.read_text(encoding="utf-8")


def register(d: Dify) -> None:
    schema = load_schema()
    body = {
        "provider": C.SERVER_STATUS_PROVIDER,
        "schema_type": "openapi",
        "schema": schema,
        "credentials": {"auth_type": "none"},
        "icon": {"background": "#2563EB", "content": "🖥"},
        "labels": ["ops", "monitoring"],
        "custom_disclaimer": "返回模拟数据，仅供 Agent 工具调用演练。",
    }
    r = d.post("/workspaces/current/tool-provider/api/add", body)
    pp("register api tool", r)
    if r.get("code") not in ("200", "201"):
        # 可能已存在，尝试 update
        body["original_provider"] = C.SERVER_STATUS_PROVIDER
        r2 = d.post("/workspaces/current/tool-provider/api/update", body)
        pp("update api tool", r2)
        if r2.get("code") not in ("200", "201"):
            raise RuntimeError("注册/更新自定义工具失败")


def list_tools(d: Dify) -> None:
    r = d.get(
        "/workspaces/current/tool-provider/api/tools?provider=%s" % C.SERVER_STATUS_PROVIDER
    )
    pp("list tools", r, limit=2000)
    r2 = d.get("/workspaces/current/tools/api")
    pp("tools/api", r2, limit=2000)


def test_tool(d: Dify) -> None:
    schema = load_schema()
    body = {
        "provider_name": C.SERVER_STATUS_PROVIDER,
        "tool_name": "getServerStatus",
        "schema_type": "openapi",
        "schema": schema,
        "credentials": {"auth_type": "none"},
        "parameters": {"server_name": "gpu-server-01"},
    }
    r = d.post("/workspaces/current/tool-provider/api/test/pre", body)
    pp("test tool", r, limit=2000)


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        register(d)
        list_tools(d)
        test_tool(d)
        return 0
    finally:
        d.close()


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