#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业 1 · 步骤 4：查 app/agent/snap/draft ID，挂工具，写提示词。"""
from __future__ import annotations

import json
import sys

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

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

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

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


def find_ids(d: Dify) -> tuple[str, str, str, str]:
    # 列表 API 返回 UTF-8 正常；不要用中文名拼进远程 psql（Windows 编码易乱）
    r_list = d.get("/agent?page=1&limit=50")
    app_id = agent_id = ""
    for item in (r_list.get("json") or {}).get("data") or []:
        if item.get("name") == C.HOMEWORK1_APP_NAME:
            app_id = item.get("id") or ""
            agent_id = item.get("agent_id") or item.get("id") or ""
            break
    if not agent_id:
        raise RuntimeError("列表里找不到应用：%s" % C.HOMEWORK1_APP_NAME)

    sql = (
        "SELECT a.app_id, a.id, a.active_config_snapshot_id, d.id "
        "FROM agents a JOIN agent_config_drafts d ON d.agent_id=a.id "
        "WHERE a.id='%s' LIMIT 1;" % agent_id
    )
    out, err = d.sh(
        'docker exec %s psql -U postgres -d dify -t -A -F "|" -c "%s"'
        % (C.PG_CONTAINER, sql)
    )
    line = (out or "").strip().splitlines()
    if not line or "|" not in line[0]:
        # 有时列表的 id 就是 app_id
        sql2 = (
            "SELECT a.app_id, a.id, a.active_config_snapshot_id, d.id "
            "FROM agents a JOIN agent_config_drafts d ON d.agent_id=a.id "
            "WHERE a.app_id='%s' LIMIT 1;" % app_id
        )
        out, err = d.sh(
            'docker exec %s psql -U postgres -d dify -t -A -F "|" -c "%s"'
            % (C.PG_CONTAINER, sql2)
        )
        line = (out or "").strip().splitlines()
    if not line or "|" not in line[0]:
        raise RuntimeError("找不到快照/草稿: %s %s" % (out, err))
    app_id, agent_id, snap_id, draft_id = line[0].split("|")
    return app_id, agent_id, snap_id, draft_id


def resolve_tool_meta(d: Dify) -> dict:
    # provider_id 必须是 tool_api_providers.id（UUID），不能用 name
    out, err = d.sh(
        "docker exec %s psql -U postgres -d dify -t -A -c "
        "\"SELECT id FROM tool_api_providers WHERE name='%s' LIMIT 1;\""
        % (C.PG_CONTAINER, C.SERVER_STATUS_PROVIDER)
    )
    provider_uuid = (out or "").strip()
    if not provider_uuid:
        raise RuntimeError("找不到 tool_api_providers.name=%s: %s" % (C.SERVER_STATUS_PROVIDER, err))

    r = d.get(
        "/workspaces/current/tool-provider/api/tools?provider=%s"
        % C.SERVER_STATUS_PROVIDER
    )
    pp("api tools meta", r, limit=2500)
    data = r.get("json") or []
    if isinstance(data, dict):
        data = data.get("data") or data.get("tools") or [data]
    if not data:
        raise RuntimeError("未找到 provider=%s 的工具列表" % C.SERVER_STATUS_PROVIDER)
    tool = data[0]
    tool_name = tool.get("name") or tool.get("tool_name") or "getServerStatus"
    return {
        "enabled": True,
        "provider_type": "api",
        "provider_id": provider_uuid,
        "tool_name": tool_name,
        "credential_type": "unauthorized",
        "credential_ref": None,
        "runtime_parameters": {},
    }


def patch_cfg(cfg: dict, tool_entry: dict) -> dict:
    prompt = cfg.setdefault("prompt", {})
    if prompt is None:
        cfg["prompt"] = {}
        prompt = cfg["prompt"]
    prompt["system_prompt"] = SYSTEM_PROMPT

    tools = cfg.get("tools")
    if not isinstance(tools, dict):
        tools = {}
        cfg["tools"] = tools
    tools["cli_tools"] = tools.get("cli_tools") or []
    tools["dify_tools"] = [tool_entry]

    model = cfg.get("model")
    if not isinstance(model, dict):
        model = {}
        cfg["model"] = model
    model["model"] = C.QWEN_MODEL
    model["plugin_id"] = "langgenius/openai"
    model["model_provider"] = C.OPENAI_PROVIDER
    model["credential_ref"] = None
    ms = model.get("model_settings")
    if not isinstance(ms, dict):
        ms = {}
        model["model_settings"] = ms
    ms["max_tokens"] = ms.get("max_tokens") or 4096
    for k in ("stop", "top_p", "temperature", "response_format", "presence_penalty", "frequency_penalty"):
        ms.setdefault(k, None)
    return cfg


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        app_id, agent_id, snap_id, draft_id = find_ids(d)
        print("app_id", app_id)
        print("agent_id", agent_id)
        print("snap_id", snap_id)
        print("draft_id", draft_id)

        tool_entry = resolve_tool_meta(d)
        print("tool_entry", json.dumps(tool_entry, ensure_ascii=False))

        snap = AC.load_cfg(d, "agent_config_snapshots", snap_id)
        draft = AC.load_cfg(d, "agent_config_drafts", draft_id)
        AC.save_cfg(d, "agent_config_snapshots", snap_id, patch_cfg(snap, tool_entry), "hw1s")
        AC.save_cfg(d, "agent_config_drafts", draft_id, patch_cfg(draft, tool_entry), "hw1d")

        # 写回 _config 提示（不自动改文件，打印给人工确认）
        print("\n请把以下常量写入 scripts_cursor/_config.py（若尚未写）：")
        print('HOMEWORK1_APP_ID = "%s"' % app_id)
        print('HOMEWORK1_AGENT_ID = "%s"' % agent_id)
        print('HOMEWORK1_SNAP_ID = "%s"' % snap_id)
        print('HOMEWORK1_DRAFT_ID = "%s"' % draft_id)
        print("\n[WARN] 改完不要点「构建/发布」；强刷后新开对话验收。")
        return 0
    finally:
        d.close()


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