#!/usr/bin/env python3
"""
从「知你客服14号」复制为「知你客服15号」：

- **工具**：与 14 号相同（平台当前全量内置工具）。
- **可持续执行**：在 LLM 节点写入 **max_tool_iterations**（默认 28），引擎在同一轮执行内允许多次
  「模型 → 工具 → 模型 → …」迭代，便于长链路干活（读文件→写文件→再校验等），而非只调一次工具就结束。
- **提示词**：强调「持续反馈、多步工具链、任务完成判定」及末行 JSON 可选字段 `task_complete` / `progress_report` 等；
  若单次无法跑完，引导用户下轮「继续」并依赖会话记忆接续。

用法:
  cd backend && .\\venv\\Scripts\\python.exe scripts/create_zhini_kefu_15.py

环境变量: PLATFORM_BASE_URL, PLATFORM_USERNAME, PLATFORM_PASSWORD,
          SOURCE_AGENT_NAME（默认 知你客服14号）, TARGET_NAME（默认 知你客服15号）
"""
from __future__ import annotations

import copy
import json
import os
import sys
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple

import requests

BASE = os.getenv("PLATFORM_BASE_URL", "http://127.0.0.1:8037").rstrip("/")
USER = os.getenv("PLATFORM_USERNAME", "admin")
PWD = os.getenv("PLATFORM_PASSWORD", "123456")
SOURCE_NAME = os.getenv("SOURCE_AGENT_NAME", "知你客服14号")
TARGET_NAME = os.getenv("TARGET_NAME", "知你客服15号")

TOOLS_V15: List[str] = [
    "http_request",
    "file_read",
    "file_write",
    "text_analyze",
    "datetime",
    "math_calculate",
    "system_info",
    "json_process",
    "database_query",
    "adb_log",
]

# 与引擎 workflow_engine 中读取的字段一致（上限 64）
# 15 号强调可持续执行，但避免过高迭代导致无效工具打转
DEFAULT_MAX_TOOL_ITERATIONS = 14

PROMPT_V15_MARKER = "【知你客服 15 号 · 可持续任务执行】"

PROMPT_V15_EXTRA = f"""

{PROMPT_V15_MARKER}

【角色】你是**可持续执行型**客服助手：面对需要多步工具配合的任务（如：查路径 → 读配置 → 写文件 → 再读回校验），应在**同一轮对话的一次执行**内，**连续使用工具**并根据返回结果决定下一步，直到任务完成或明确受阻；不要只做一次工具调用就结束。

【与 14 号的关系】继承 14 号全部内置工具与纪律；**工具列表未删减**，平台侧已为 15 号提高**单次执行内工具迭代次数**（见节点 `max_tool_iterations`）。

【执行策略】
1. **默认本地闭环**：先 `system_info` 确认工作区，再 `file_read/file_write/text_analyze` 完成本地任务；仅当用户**明确要求联网检索**（如“上网查”“联网获取”）时才可调用 `http_request`。
2. **持续反馈**：在最终自然语言中说明**已做步骤**与**当前结果**；勿编造工具返回。
3. **何时停**：目标达成 → 在末行 JSON 中标明完成；缺用户输入/权限/环境 → 清楚说明缺什么。
4. **单次装不下时**：在 `reply` 中说明进度，并建议用户**下一轮发送「继续」**；可把未完成要点写入 `user_profile` 或依赖会话记忆中的 `conversation_history` 衔接（勿用空 JSON 覆盖画像）。
5. **古文/常识续写类任务**（如《三字经》补全段落）：视为通用知识，不得为此调用 `http_request`；应直接给出内容并按需落盘。

【末行 JSON（单行）扩展字段（推荐）】
在原有 `intent`、`reply`、`user_profile` 基础上，可增加：
- `task_complete`: boolean，本任务是否已彻底完成；
- `progress_report`: string，本轮已完成步骤的简要清单；
- `continuation_hint`: string，若 `task_complete` 为 false，提示用户下一句怎么说（如「继续」「补充 xxx」）。

仍须以 **一行合法 JSON** 结尾，勿用 markdown 代码围栏。

【纪律】继承 14 号：勿刷屏 DSML；严禁把 `<｜DSML｜...>`、工具调用协议原文输出给用户；`database_query` 仅 SELECT；`file_write` 同轮勿无故重复写入同一文件除非必要。
"""


def _sanitize_edges(edges: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    seen: set = set()
    out: List[Dict[str, Any]] = []
    for e in edges or []:
        s, t = e.get("source"), e.get("target")
        if not s or not t:
            continue
        if s == t:
            continue
        key = (s, t)
        if key in seen:
            continue
        seen.add(key)
        ne = dict(e)
        ne["sourceHandle"] = "right"
        ne["targetHandle"] = "left"
        if not ne.get("id"):
            ne["id"] = f"edge_{s}_{t}"
        out.append(ne)
    return out


def _find_start_node_ids(nodes: List[Dict[str, Any]]) -> List[str]:
    ids: List[str] = []
    for n in nodes or []:
        nid = n.get("id") or ""
        nt = (n.get("type") or (n.get("data") or {}).get("type") or "").lower()
        if nt == "start" or nid in ("start", "start-1") or str(nid).startswith("start-"):
            ids.append(nid)
    return ids


def _compute_ranks(
    nodes: List[Dict[str, Any]], edges: List[Dict[str, Any]]
) -> Dict[str, int]:
    node_ids = [n["id"] for n in nodes if n.get("id")]
    start_ids = _find_start_node_ids(nodes)
    incoming: Dict[str, int] = {nid: 0 for nid in node_ids}
    for e in edges:
        s, t = e.get("source"), e.get("target")
        if not s or not t or s == t:
            continue
        if t in incoming:
            incoming[t] += 1
    if not start_ids:
        start_ids = [nid for nid in node_ids if incoming.get(nid, 0) == 0] or ([node_ids[0]] if node_ids else [])

    rank: Dict[str, int] = {s: 0 for s in start_ids}
    nmax = max(len(nodes), 8)
    for _ in range(nmax + 5):
        updated = False
        for e in edges:
            s, t = e.get("source"), e.get("target")
            if not s or not t or s == t:
                continue
            if s not in rank:
                continue
            nv = rank[s] + 1
            if t not in rank or rank[t] < nv:
                rank[t] = nv
                updated = True
        if not updated:
            break
    max_r = max(rank.values(), default=0)
    for nid in node_ids:
        if nid not in rank:
            rank[nid] = max_r + 1
            max_r += 1
    return rank


def _apply_layered_positions(nodes: List[Dict[str, Any]], ranks: Dict[str, int]) -> None:
    layers: Dict[int, List[str]] = defaultdict(list)
    for nid, r in ranks.items():
        layers[r].append(nid)
    for r in layers:
        layers[r].sort()

    x0, y0 = 80.0, 140.0
    x_step = 300.0
    y_step = 110.0

    for r in sorted(layers.keys()):
        ids = layers[r]
        nlen = len(ids)
        y_base = y0 - (nlen - 1) * y_step / 2.0
        for j, nid in enumerate(ids):
            for node in nodes:
                if node.get("id") != nid:
                    continue
                pos = node.setdefault("position", {})
                pos["x"] = x0 + r * x_step
                pos["y"] = y_base + j * y_step
                break


def improve_workflow_layout_and_edges(wf: Dict[str, Any]) -> Tuple[int, int]:
    nodes = wf.get("nodes") or []
    raw_edges = wf.get("edges") or []
    loops = sum(
        1
        for e in raw_edges
        if e.get("source") and e.get("target") and e.get("source") == e.get("target")
    )
    clean = _sanitize_edges(raw_edges)
    removed_dup = len(raw_edges) - len(clean) - loops

    wf["edges"] = clean

    ranks = _compute_ranks(nodes, clean)
    _apply_layered_positions(nodes, ranks)
    return loops, max(0, removed_dup)


def _patch_llm_unified(wf: dict, base_prompt: Optional[str] = None) -> None:
    for n in wf.get("nodes") or []:
        if n.get("id") != "llm-unified":
            continue
        d = n.setdefault("data", {})
        prompt = base_prompt if base_prompt else d.get("prompt") or ""
        if PROMPT_V15_MARKER not in prompt:
            prompt = (prompt.rstrip() + "\n" + PROMPT_V15_EXTRA).strip()
        d["prompt"] = prompt
        d["enable_tools"] = True
        d["tools"] = list(TOOLS_V15)
        d["selected_tools"] = list(TOOLS_V15)
        d["max_tool_iterations"] = DEFAULT_MAX_TOOL_ITERATIONS
        return
    print("警告: 未找到节点 llm-unified", file=sys.stderr)


def _find_agent_id_by_name(h: Dict[str, str], name: str) -> Optional[str]:
    r = requests.get(f"{BASE}/api/v1/agents", params={"search": name, "limit": 50}, headers=h, timeout=30)
    if r.status_code != 200:
        return None
    for a in r.json() or []:
        if a.get("name") == name:
            return a.get("id")
    return None


def main() -> int:
    r = requests.post(
        f"{BASE}/api/v1/auth/login",
        data={"username": USER, "password": PWD},
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        timeout=15,
    )
    if r.status_code != 200:
        print("登录失败:", r.status_code, r.text[:500], file=sys.stderr)
        return 1
    token = r.json().get("access_token")
    if not token:
        print("无 access_token", file=sys.stderr)
        return 1
    h = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

    src_id = _find_agent_id_by_name(h, SOURCE_NAME)
    if not src_id:
        print(f"未找到源 Agent: {SOURCE_NAME}", file=sys.stderr)
        return 1

    existing = _find_agent_id_by_name(h, TARGET_NAME)
    if existing:
        print("已存在", TARGET_NAME, "-> 仅更新工作流", existing)
        new_id = existing
        g = requests.get(f"{BASE}/api/v1/agents/{new_id}", headers=h, timeout=30)
        if g.status_code != 200:
            print("读取失败:", g.text, file=sys.stderr)
            return 1
        agent = g.json()
    else:
        dup = requests.post(
            f"{BASE}/api/v1/agents/{src_id}/duplicate",
            headers=h,
            json={"name": TARGET_NAME},
            timeout=60,
        )
        if dup.status_code != 201:
            print("复制失败:", dup.status_code, dup.text[:800], file=sys.stderr)
            return 1
        new_id = dup.json()["id"]
        agent = dup.json()
        print("已创建副本:", new_id, TARGET_NAME)

    wf = copy.deepcopy(agent["workflow_config"])
    loops, dup_edges = improve_workflow_layout_and_edges(wf)
    print(f"连线整理: 去掉自环 {loops} 条, 合并重复边 {dup_edges} 条")

    g2 = requests.get(f"{BASE}/api/v1/agents/{src_id}", headers=h, timeout=30)
    base_prompt = None
    if g2.status_code == 200:
        try:
            for n in g2.json().get("workflow_config", {}).get("nodes") or []:
                if n.get("id") == "llm-unified":
                    base_prompt = (n.get("data") or {}).get("prompt")
                    break
        except Exception:
            pass
    _patch_llm_unified(wf, base_prompt=base_prompt)

    desc = (
        "知你客服15号：在14号全量工具基础上，强调可持续多步执行；"
        f"llm-unified 配置 max_tool_iterations={DEFAULT_MAX_TOOL_ITERATIONS}，"
        "单次执行内可多轮工具调用直至任务完成或明确需用户继续；输出单行 JSON，可含 task_complete/progress_report。"
    )

    up = requests.put(
        f"{BASE}/api/v1/agents/{new_id}",
        headers=h,
        json={"description": desc, "workflow_config": wf},
        timeout=120,
    )
    if up.status_code != 200:
        print("更新失败:", up.status_code, up.text[:1200], file=sys.stderr)
        return 1
    print("已写入工具:", ", ".join(TOOLS_V15))
    print(f"max_tool_iterations: {DEFAULT_MAX_TOOL_ITERATIONS}")
    print("Agent ID:", new_id)
    print(json.dumps({"id": new_id, "name": TARGET_NAME}, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())