#!/usr/bin/env python3
"""在现有 Dify 上跑三条工作流：工单分支、监控取数、客服检索。

不新建模型，不重启 vLLM。同名应用已存在时覆盖草稿再发布。
每次只跑一条。显存里大约 86GB 是 vLLM 预先划出的缓存。
"""
from __future__ import annotations

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

BASE = "http://127.0.0.1:8088/console/api"
PROVIDER = "langgenius/openai_api_compatible/openai_api_compatible"
RERANK = "langgenius/xinference/xinference"
FAQ_DATASET = "b2dd0af3-d73c-4983-9348-8b25ca502ef2"
MODEL = {
    "provider": PROVIDER,
    "name": "qwen3.6-27b",
    "mode": "chat",
    "completion_params": {"temperature": 0.2, "max_tokens": 2048},
}
FEATURES = {
    "opening_statement": "",
    "suggested_questions": [],
    "suggested_questions_after_answer": {"enabled": False},
    "text_to_speech": {"enabled": False, "voice": "", "language": ""},
    "speech_to_text": {"enabled": False},
    "retriever_resource": {"enabled": True},
    "sensitive_word_avoidance": {"enabled": False},
    "file_upload": {
        "enabled": False,
        "allowed_file_types": ["image"],
        "allowed_file_extensions": [],
        "allowed_file_upload_methods": ["local_file", "remote_url"],
        "number_limits": 3,
        "image": {
            "enabled": False,
            "number_limits": 3,
            "transfer_methods": ["local_file", "remote_url"],
        },
    },
}


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[:800]}
            return exc.code, body


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 node(nid: str, data: dict, x: int, y: int) -> dict:
    return {
        "id": nid,
        "type": "custom",
        "data": data,
        "position": {"x": x, "y": y},
        "positionAbsolute": {"x": x, "y": y},
        "width": 244,
        "height": 90,
        "sourcePosition": "right",
        "targetPosition": "left",
    }


def edge(source: str, target: str, source_type: str, target_type: str, handle: str = "source") -> dict:
    return {
        "id": "%s-%s-%s" % (source, handle, target),
        "type": "custom",
        "source": source,
        "target": target,
        "sourceHandle": handle,
        "targetHandle": "target",
        "data": {
            "sourceType": source_type,
            "targetType": target_type,
            "isInIteration": False,
            "isInLoop": False,
        },
        "zIndex": 0,
    }


def start_node(variables: list[dict], y: int = 200) -> dict:
    return node(
        "start",
        {"type": "start", "title": "开始", "variables": variables},
        80,
        y,
    )


def end_node(nid: str, source: str, x: int, y: int, output_key: str = "output") -> dict:
    return node(
        nid,
        {
            "type": "end",
            "title": "结束",
            "outputs": [
                {
                    "variable": "result",
                    "value_type": "string",
                    "value_selector": [source, output_key],
                }
            ],
        },
        x,
        y,
    )


def http_node(nid: str, title: str, query: str, x: int, y: int) -> dict:
    return node(
        nid,
        {
            "type": "http-request",
            "title": title,
            "method": "get",
            "url": "http://10.0.1.47:9090/api/v1/query",
            "authorization": {"type": "no-auth"},
            "headers": "",
            "params": "query:%s" % query,
            "body": {"type": "none", "data": []},
            "timeout": {"connect": 10, "read": 20, "write": 10},
            "ssl_verify": False,
        },
        x,
        y,
    )


def ticket_graph() -> dict:
    classifier = node(
        "cls",
        {
            "type": "question-classifier",
            "title": "判断类别",
            "query_variable_selector": ["start", "ticket"],
            "model": MODEL,
            "instruction": "催发票、抱怨、要求马上处理，选投诉。询问地址、入口或用法，选咨询。",
            "classes": [
                {"id": "complaint", "name": "投诉", "label": "投诉"},
                {"id": "consult", "name": "咨询", "label": "咨询"},
            ],
            "vision": {"enabled": False},
        },
        380,
        200,
    )
    high = node(
        "high",
        {
            "type": "template-transform",
            "title": "投诉出口",
            "variables": [{"variable": "label", "value_selector": ["cls", "class_label"]}],
            "template": "类别：{{ label }}\n紧急程度：高\n处理：今日内答复",
        },
        700,
        80,
    )
    low = node(
        "low",
        {
            "type": "template-transform",
            "title": "咨询出口",
            "variables": [{"variable": "label", "value_selector": ["cls", "class_label"]}],
            "template": "类别：{{ label }}\n紧急程度：普通\n处理：按平台说明答复",
        },
        700,
        320,
    )
    return {
        "nodes": [
            start_node(
                [
                    {
                        "variable": "ticket",
                        "label": "工单内容",
                        "type": "paragraph",
                        "required": True,
                        "max_length": 500,
                        "options": [],
                    }
                ]
            ),
            classifier,
            high,
            low,
            end_node("end_high", "high", 980, 80),
            end_node("end_low", "low", 980, 320),
        ],
        "edges": [
            edge("start", "cls", "start", "question-classifier"),
            edge("cls", "high", "question-classifier", "template-transform", "complaint"),
            edge("cls", "low", "question-classifier", "template-transform", "consult"),
            edge("high", "end_high", "template-transform", "end"),
            edge("low", "end_low", "template-transform", "end"),
        ],
        "viewport": {"x": 0, "y": 0, "zoom": 0.8},
    }


def monitor_graph() -> dict:
    code = """
def main(temp: str, mem: str, queue: str) -> dict:
    import json

    def pick(raw: str) -> str:
        try:
            data = json.loads(raw)
        except Exception:
            return (raw or "")[:160]
        rows = ((data.get("data") or {}).get("result")) or []
        if not rows:
            return "无数据"
        value = rows[0].get("value") or []
        if len(value) < 2:
            return "无数据"
        return str(value[1])

    return {"text": "温度 %s C\\n显存 %s MiB\\n排队 %s" % (pick(temp), pick(mem), pick(queue))}
"""
    return {
        "nodes": [
            start_node([]),
            http_node("http_temp", "查温度", "DCGM_FI_DEV_GPU_TEMP", 360, 200),
            http_node("http_mem", "查显存", "DCGM_FI_DEV_FB_USED", 640, 200),
            http_node(
                "http_queue",
                "查排队",
                'sum(vllm:num_requests_waiting{model_name="qwen3.6-27b"})',
                920,
                200,
            ),
            node(
                "parse",
                {
                    "type": "code",
                    "title": "抽出数值",
                    "variables": [
                        {"variable": "temp", "value_selector": ["http_temp", "body"]},
                        {"variable": "mem", "value_selector": ["http_mem", "body"]},
                        {"variable": "queue", "value_selector": ["http_queue", "body"]},
                    ],
                    "code_language": "python3",
                    "code": code,
                    "outputs": {"text": {"type": "string", "children": None}},
                },
                1200,
                200,
            ),
            end_node("end", "parse", 1480, 200, "text"),
        ],
        "edges": [
            edge("start", "http_temp", "start", "http-request"),
            edge("http_temp", "http_mem", "http-request", "http-request"),
            edge("http_mem", "http_queue", "http-request", "http-request"),
            edge("http_queue", "parse", "http-request", "code"),
            edge("parse", "end", "code", "end"),
        ],
        "viewport": {"x": 0, "y": 0, "zoom": 0.7},
    }


def support_graph() -> dict:
    return {
        "nodes": [
            start_node(
                [
                    {
                        "variable": "question",
                        "label": "问题",
                        "type": "paragraph",
                        "required": True,
                        "max_length": 300,
                        "options": [],
                    }
                ]
            ),
            node(
                "kr",
                {
                    "type": "knowledge-retrieval",
                    "title": "检索平台说明",
                    "query_variable_selector": ["start", "question"],
                    "dataset_ids": [FAQ_DATASET],
                    "retrieval_mode": "multiple",
                    "multiple_retrieval_config": {
                        "top_k": 3,
                        "reranking_enable": True,
                        "reranking_mode": "reranking_model",
                        "reranking_model": {
                            "reranking_provider_name": RERANK,
                            "reranking_model_name": "bge-reranker-v2-m3",
                        },
                    },
                    "metadata_filtering_mode": "disabled",
                },
                380,
                200,
            ),
            node(
                "llm",
                {
                    "type": "llm",
                    "title": "根据说明回答",
                    "model": MODEL,
                    "prompt_template": [
                        {
                            "role": "system",
                            "text": "下面是检索到的平台说明。只根据说明回答。说明里没有的就说不知道。不要编造口令。最多两句。\n{{#context#}}",
                            "edition_type": "basic",
                        },
                        {
                            "role": "user",
                            "text": "{{#start.question#}}",
                            "edition_type": "basic",
                        },
                    ],
                    "context": {"enabled": True, "variable_selector": ["kr", "result"]},
                    "vision": {"enabled": False},
                    "reasoning_format": "separated",
                },
                700,
                200,
            ),
            end_node("end", "llm", 980, 200, "text"),
        ],
        "edges": [
            edge("start", "kr", "start", "knowledge-retrieval"),
            edge("kr", "llm", "knowledge-retrieval", "llm"),
            edge("llm", "end", "llm", "end"),
        ],
        "viewport": {"x": 0, "y": 0, "zoom": 0.8},
    }


def ensure(dify: Dify, name: str, description: str, graph: dict) -> str:
    _status, listed = dify.call("GET", "/apps?page=1&limit=50&mode=workflow")
    app_id = ""
    for item in listed.get("data") or []:
        if item.get("name") == name:
            app_id = item["id"]
            break
    if not app_id:
        status, created = dify.call(
            "POST",
            "/apps",
            {
                "name": name,
                "description": description,
                "mode": "workflow",
                "icon_type": "emoji",
                "icon": "🔀",
                "icon_background": "#E3F2FD",
            },
        )
        app_id = created.get("id") if isinstance(created, dict) else ""
        if not app_id:
            raise SystemExit("创建工作流失败 %s %s" % (status, created))
        print("已创建", name, app_id)
    else:
        print("工作流已在", name, app_id)
    _status, draft = dify.call("GET", "/apps/%s/workflows/draft" % app_id)
    current_hash = draft.get("hash") if isinstance(draft, dict) else None
    status, synced = dify.call(
        "POST",
        "/apps/%s/workflows/draft" % app_id,
        {
            "graph": graph,
            "features": FEATURES,
            "hash": current_hash,
            "conversation_variables": [],
        },
    )
    if status >= 400:
        raise SystemExit("保存草稿失败 %s %s" % (status, synced))
    status, published = dify.call("POST", "/apps/%s/workflows/publish" % app_id, {})
    print("已发布", name, status, published.get("result") if isinstance(published, dict) else published)
    if status >= 400:
        raise SystemExit(published)
    return app_id


def run(dify: Dify, app_id: str, inputs: dict) -> str:
    print("显存占用", gpu_used(), "MiB")
    status, raw = dify.call(
        "POST",
        "/apps/%s/workflows/draft/run" % app_id,
        {"inputs": inputs},
        timeout=180,
    )
    if status != 200:
        return "HTTP %s %s" % (status, raw if isinstance(raw, dict) else str(raw)[:400])
    text = raw if isinstance(raw, str) else json.dumps(raw, ensure_ascii=False)
    steps = []
    outputs = None
    failed = ""
    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")
        data = event.get("data") or {}
        if kind == "node_finished":
            title = data.get("title") or data.get("node_id")
            node_status = data.get("status")
            steps.append("%s=%s" % (title, node_status))
            if data.get("error"):
                failed = str(data.get("error"))[:240]
        if kind == "workflow_finished":
            outputs = data.get("outputs")
            if data.get("status") not in (None, "succeeded"):
                failed = failed or str(data.get("error") or data.get("status"))[:240]
        if kind == "error":
            failed = str(event.get("message") or event)[:240]
    body = json.dumps(outputs, ensure_ascii=False) if outputs is not None else ""
    summary = "步骤 %s" % " -> ".join(steps)
    if failed:
        summary += "\n失败 " + failed
    if body:
        summary += "\n输出 " + body[:500]
    return summary or text[:400]


def main() -> None:
    dify = Dify()
    dify.login()
    ticket = ensure(dify, "工单分流工作流", "先分类，再走投诉或咨询出口", ticket_graph())
    monitor = ensure(dify, "监控检查工作流", "按固定顺序读取温度、显存和排队", monitor_graph())
    support = ensure(dify, "客服检索工作流", "先检索平台说明，再回答", support_graph())

    print("==== 工单分流工作流 / 投诉 ====")
    print(run(dify, ticket, {"ticket": "合同发票已经三天没收到，请今天给一个答复"}))
    print("==== 工单分流工作流 / 咨询 ====")
    print(run(dify, ticket, {"ticket": "GitLab 的网页地址在哪里"}))
    print("==== 监控检查工作流 ====")
    print(run(dify, monitor, {}))
    print("==== 客服检索工作流 ====")
    print(run(dify, support, {"question": "GitLab 的网页地址是什么"}))
    print("结束显存", gpu_used(), "MiB")
    health = request.urlopen("http://127.0.0.1:8000/health", timeout=10)
    print("vllm", health.status)
    print("WORKFLOWS_DONE")


if __name__ == "__main__":
    main()