#!/usr/bin/env python3
"""在现有 Dify 上试用客服、工单、监控和只读自愈。

不新建模型，不重启 vLLM。同名应用和知识库已存在时只复用。
每次只发一条短对话。nvidia-smi 里大约 86GB 是 vLLM 启动时划定的缓存，不是满负荷。
"""
from __future__ import annotations

import base64
import json
import subprocess
import time
from http.cookiejar import CookieJar
from pathlib import Path
from urllib import error, request

BASE = "http://127.0.0.1:8088/console/api"
ROOT = Path(__file__).resolve().parents[1]
PROVIDER = "langgenius/xinference/xinference"
CONTRACT_AGENT = "01a1236c-4cc3-7fb9-bfce-c730a318bdf9"
OPS_APP = "e4d614f6-de22-4f4c-b58e-caeb0504b121"
FAQ_NAME = "平台客服说明"
FAQ_TEXT = """平台入口说明

Dify 网页是 http://10.0.1.47:8088 ，用来做对话、知识库和工具调用。
GitLab 网页是 http://10.0.1.47:8929 ，私有仓库是 gitops/platform，登录用户名是 root。口令写在导航页，客服不要编造口令。
Argo CD 网页是 http://10.0.1.47:8080 ，登录用户名是 admin。它负责把 Git 清单同步到集群。
对话模型名称是 qwen3.6-27b，请求从 Dify 到 New-API，再到 vLLM。
知识库向量模型是 bge-m3，重排模型是 bge-reranker-v2-m3，由 Xinference 提供。
忘记地址时，先看导航页 http://10.0.1.47:8090 。
"""


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, timeout: int = 60):
        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)
        try:
            with self.opener.open(req, timeout=timeout) as resp:
                raw = resp.read().decode()
                if not raw:
                    return resp.status, {}
                try:
                    return resp.status, json.loads(raw)
                except json.JSONDecodeError:
                    return resp.status, raw
        except error.HTTPError as exc:
            raw = exc.read().decode(errors="replace")
            try:
                body = json.loads(raw)
            except json.JSONDecodeError:
                body = {"err": raw[:500]}
            return exc.code, body

    def upload_text(self, filename: str, content: str) -> str:
        boundary = "----difyboundary7"
        head = (
            "--%s\r\n"
            'Content-Disposition: form-data; name="file"; filename="%s"\r\n'
            "Content-Type: text/plain\r\n\r\n" % (boundary, filename)
        ).encode()
        tail = (
            "\r\n--%s\r\n"
            'Content-Disposition: form-data; name="source"\r\n\r\n'
            "datasets\r\n"
            "--%s--\r\n" % (boundary, boundary)
        ).encode()
        req = request.Request(BASE + "/files/upload", data=head + content.encode() + tail, method="POST")
        req.add_header("Content-Type", "multipart/form-data; boundary=%s" % boundary)
        req.add_header("X-CSRF-Token", self.csrf)
        try:
            with self.opener.open(req, timeout=60) as resp:
                body = json.loads(resp.read().decode())
        except error.HTTPError as exc:
            raise SystemExit("上传说明失败 %s %s" % (exc.code, exc.read().decode(errors="replace")[:400]))
        file_id = body.get("id")
        if not file_id:
            raise SystemExit("上传没有返回文件编号 %s" % body)
        return file_id


def gpu_used() -> int:
    out = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
        text=True,
    )
    return int(out.strip().splitlines()[0])


def enable_trace(dify: Dify, app_id: str) -> None:
    env_path = Path("/usr/cluster/observability/langfuse/.env")
    if not env_path.exists():
        print("跳过追踪，没有 langfuse env")
        return
    env: dict[str, str] = {}
    for line in env_path.read_text(encoding="utf-8").splitlines():
        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"},
    )


def ensure_tool(dify: Dify, provider: str, schema_path: Path) -> str:
    _status, tools = dify.call("GET", "/workspaces/current/tools/api")
    rows = tools if isinstance(tools, list) else (tools.get("data") or [])
    for item in rows:
        if item.get("name") == provider or item.get("provider") == provider:
            print("工具已在", provider, item.get("id"))
            return item["id"]
    status, created = dify.call(
        "POST",
        "/workspaces/current/tool-provider/api/add",
        {
            "provider": provider,
            "credentials": {"auth_type": "none"},
            "schema_type": "openapi",
            "schema": schema_path.read_text(encoding="utf-8"),
            "icon": {"background": "#E3F2FD", "content": "P"},
            "privacy_policy": "",
            "custom_disclaimer": "",
            "labels": [],
        },
    )
    print("注册工具", provider, status)
    _status, tools = dify.call("GET", "/workspaces/current/tools/api")
    rows = tools if isinstance(tools, list) else (tools.get("data") or [])
    for item in rows:
        if item.get("name") == provider or item.get("provider") == provider:
            return item["id"]
    raise SystemExit("找不到工具编号 %s %s" % (provider, created))


def tool_entry(provider_id: str, tool_name: str) -> dict:
    return {
        "enabled": True,
        "provider_type": "api",
        "provider_id": provider_id,
        "tool_name": tool_name,
        "credential_type": "unauthorized",
        "credential_ref": None,
        "runtime_parameters": {},
    }


def ensure_dataset(dify: Dify) -> str:
    _status, listed = dify.call("GET", "/datasets?page=1&limit=50")
    dataset_id = ""
    for item in listed.get("data") or []:
        if item.get("name") == FAQ_NAME:
            dataset_id = item["id"]
            break
    _status, sample = dify.call("GET", "/datasets/71042108-5f9a-4856-a0db-4359e8e7fa98")
    retrieval = sample.get("retrieval_model_dict") or sample.get("retrieval_model")
    if not dataset_id:
        status, created = dify.call(
            "POST",
            "/datasets",
            {
                "name": FAQ_NAME,
                "description": "内部平台入口，供客服检索",
                "indexing_technique": "high_quality",
                "permission": "all_team_members",
                "embedding_model": "bge-m3",
                "embedding_model_provider": PROVIDER,
                "retrieval_model": retrieval,
            },
        )
        dataset_id = created.get("id")
        if not dataset_id:
            raise SystemExit("创建知识库失败 %s %s" % (status, created))
        print("已创建知识库", dataset_id)
    else:
        print("知识库已在", dataset_id)
    _status, docs = dify.call("GET", "/datasets/%s/documents?page=1&limit=5" % dataset_id)
    rows = docs.get("data") or []
    if rows and rows[0].get("indexing_status") == "completed":
        print("说明已索引")
        return dataset_id
    file_id = dify.upload_text("platform-faq.txt", FAQ_TEXT)
    print("已上传", file_id)
    status, doc = dify.call(
        "POST",
        "/datasets/%s/documents" % dataset_id,
        {
            "name": "平台入口",
            "indexing_technique": "high_quality",
            "doc_form": "text_model",
            "doc_language": "Chinese",
            "embedding_model": "bge-m3",
            "embedding_model_provider": PROVIDER,
            "data_source": {
                "info_list": {
                    "data_source_type": "upload_file",
                    "file_info_list": {"file_ids": [file_id]},
                }
            },
            "process_rule": {"mode": "automatic"},
            "retrieval_model": retrieval,
        },
        timeout=120,
    )
    print("写入说明", status)
    if status >= 400:
        raise SystemExit(doc)
    for _ in range(40):
        _status, docs = dify.call("GET", "/datasets/%s/documents?page=1&limit=5" % dataset_id)
        rows = docs.get("data") or []
        indexing = rows[0].get("indexing_status") or rows[0].get("display_status") or "empty"
        if _ == 0:
            print("文档字段", sorted(rows[0].keys())[:12])
        print("索引", indexing)
        if indexing == "completed":
            return dataset_id
        if indexing == "error":
            raise SystemExit("客服说明索引失败 %s" % rows[0].get("error"))
        time.sleep(3)
    raise SystemExit("客服说明索引超时")


def knowledge_set(dataset_id: str) -> dict:
    return {
        "id": "retrieval-1",
        "name": "Retrieval 1",
        "description": "平台入口说明",
        "datasets": [{"id": dataset_id, "name": FAQ_NAME}],
        "query": {"mode": "generated_query"},
        "retrieval": {
            "mode": "multiple",
            "top_k": 3,
            "reranking_mode": "reranking_model",
            "reranking_enable": True,
            "reranking_model": {"provider": PROVIDER, "model": "bge-reranker-v2-m3"},
        },
        "metadata_filtering": {"mode": "disabled"},
    }


def ensure_agent(dify: Dify, name: str, description: str, prompt: str, tools: list, dataset_id: str | None) -> str:
    _status, listed = dify.call("GET", "/agent?page=1&limit=50")
    for item in listed.get("data") or []:
        if item.get("name") == name:
            app_id = item.get("app_id") or item.get("backing_app_id")
            print("应用已在", name, app_id)
            return app_id
    _status, state = dify.call("GET", "/agent/%s/composer" % CONTRACT_AGENT)
    model = dict(state["agent_soul"]["model"])
    params = dict(model.get("completion_params") or {})
    params["max_tokens"] = 256
    params["temperature"] = 0.2
    model["completion_params"] = params
    status, created = dify.call(
        "POST",
        "/agent",
        {
            "name": name,
            "description": description,
            "icon_type": "emoji",
            "icon": "💬",
            "icon_background": "#E3F2FD",
            "role": "assistant",
        },
    )
    agent_id = created.get("agent_id") or created.get("id")
    if not agent_id:
        raise SystemExit("创建应用失败 %s %s" % (status, created))
    _status, state = dify.call("GET", "/agent/%s/composer" % agent_id)
    soul = state["agent_soul"]
    soul.setdefault("prompt", {})["system_prompt"] = prompt
    soul.setdefault("tools", {})["dify_tools"] = tools
    soul["tools"]["cli_tools"] = soul["tools"].get("cli_tools") or []
    soul["model"] = model
    if dataset_id:
        soul.setdefault("knowledge", {})["sets"] = [knowledge_set(dataset_id)]
    dify.call(
        "PUT",
        "/agent/%s/composer" % agent_id,
        {
            "variant": "agent_app",
            "save_strategy": "save_to_current_version",
            "agent_soul": soul,
            "version_note": name,
        },
    )
    dify.call("POST", "/agent/%s/publish" % agent_id, {"version_note": name})
    app_id = state.get("app_id") or state.get("backing_app_id")
    print("已发布", name, app_id)
    if app_id:
        enable_trace(dify, app_id)
    return app_id


def chat(dify: Dify, app_id: str, query: str) -> str:
    print("显存占用", gpu_used(), "MiB，利用率见 nvidia-smi")
    status, raw = dify.call(
        "POST",
        "/apps/%s/chat-messages" % app_id,
        {
            "query": query,
            "inputs": {},
            "response_mode": "streaming",
            "conversation_id": "",
            "user": "scenario",
        },
        timeout=180,
    )
    if status != 200:
        return "HTTP %s %s" % (status, raw if isinstance(raw, dict) else raw)
    text = raw if isinstance(raw, str) else json.dumps(raw, ensure_ascii=False)
    answer = []
    tools = []
    for line in text.splitlines():
        if not line.startswith("data:"):
            continue
        try:
            event = json.loads(line[5:].strip())
        except json.JSONDecodeError:
            continue
        kind = event.get("event")
        if kind in ("agent_message", "message") and event.get("answer"):
            answer.append(event["answer"])
        if kind == "agent_thought" and event.get("tool"):
            tools.append(event["tool"])
        if kind == "error":
            return "ERROR " + str(event.get("message") or event)[:300]
    body = "".join(answer).strip()
    if tools:
        body = "[工具 %s] %s" % (",".join(tools), body)
    return body[:500] or text[:300]


def main() -> None:
    dify = Dify()
    dify.login()
    prom_id = ensure_tool(dify, "prometheus_query", ROOT / "config" / "prometheus-openapi.yaml")
    health_id = ensure_tool(dify, "vllm_health", ROOT / "config" / "vllm-health-openapi.yaml")
    dataset_id = ensure_dataset(dify)

    service = ensure_agent(
        dify,
        "智能客服",
        "根据平台说明回答入口和用法",
        "你是内部平台客服。只能根据知识库回答。不知道就说不知道，不要编造口令。回答不超过 80 字。",
        [],
        dataset_id,
    )
    ticket = ensure_agent(
        dify,
        "工单分流",
        "把来话分成咨询、投诉、故障、账号",
        "你是客服工单分流员。把用户来话分成咨询、投诉、故障、账号之一，并给出紧急程度高或低。用三行输出：类别、紧急程度、摘要。不要调用工具。",
        [],
        None,
    )
    monitor = ensure_agent(
        dify,
        "监控检查",
        "用 Prometheus 读取 GPU 和 vLLM 排队",
        """你是监控值班员。必须调用 queryPrometheus，不要编造数字。
依次查询这三条，并各用一行汇报数值：
DCGM_FI_DEV_GPU_TEMP
DCGM_FI_DEV_FB_USED
sum(vllm:num_requests_waiting{model_name="qwen3.6-27b"})
温度单位沿用指标原值，显存沿用指标原值。回答不超过 120 字。""",
        [tool_entry(prom_id, "queryPrometheus")],
        None,
    )
    heal = ensure_agent(
        dify,
        "故障自愈",
        "只读检查 vLLM 健康和排队，给出是否处理",
        """你是只读的故障自愈值班员。必须调用 checkVllmHealth 和 queryPrometheus。
Prometheus 查询 sum(vllm:num_requests_waiting{model_name="qwen3.6-27b"})。
结论只写：无需处理、观察、人工处理 三者之一。
排队大于 10 时写观察，并说明 KEDA 会把 Dify worker 扩到最多 3 个。
健康检查返回空内容、且调用没有报错，就当作成功。
健康检查调用失败时写人工处理。
不要输出重启、删除 Pod、kubectl 或改镜像的命令。回答不超过 100 字。""",
        [tool_entry(health_id, "checkVllmHealth"), tool_entry(prom_id, "queryPrometheus")],
        None,
    )

    cases = [
        ("工单分流", ticket, "客户说合同发票三天了还没收到，语气很急，要求今天给答复。"),
        ("监控检查", monitor, "检查当前 GPU 温度、显存占用和 vLLM 排队人数。"),
        ("故障自愈", heal, "判断现在要不要处理 vLLM。只给结论，不要执行任何修复。"),
        ("智能客服", service, "GitLab 的网页地址是什么？用一句话回答。"),
        ("合同问答", "d2e3f57b-3901-45de-bd9a-d51feb849f6e", "这份合同的客户是谁？只根据知识库，用一句话回答。"),
        ("运维巡检", OPS_APP, "查询服务器 gpu-server-01 的状态，按工具结果简要汇报。"),
    ]
    for name, app_id, query in cases:
        print("====", name, "====")
        if not app_id:
            print("没有应用编号")
            continue
        print(chat(dify, app_id, query))
    health = subprocess.check_output(
        ["curl", "-s", "-o", "/dev/null", "-w", "%{http_code}", "--max-time", "5", "http://127.0.0.1:8000/health"],
        text=True,
    )
    print("结束显存", gpu_used(), "MiB", "vllm", health)
    print("SCENARIOS_DONE")


if __name__ == "__main__":
    main()