#!/usr/bin/env python3
"""
批量创建多个「企业场景」示例 Agent（画布均为无环 DAG，可通过平台校验）。

场景一览：
  1. 企业场景_电商售后       — LLM + http / 文本分析 / 时间
  2. 企业场景_研发联调       — LLM + http / JSON / 读文件 / 数学
  3. 企业场景_数据分析       — LLM + 只读库查询 / JSON / 文本 / 数学（需数据源与权限）
  4. 企业场景_移动运维       — LLM + adb / 读文件 / 文本 / 系统信息
  5. 企业场景_合规纪要       — 纯 LLM，偏合规与留痕表述
  6. 企业场景_销售助理       — LLM + http / text_analyze / datetime
  7. 企业场景_多线路由       — Code 解析 query 中的线路标记 → Switch → 三条专精 LLM → Merge
  8. 企业场景_审批流样例     — Approval → 通过则 LLM 生成说明 / 拒绝则直接结束（首次执行会待审批）

多线路由标记（写在用户 query 里，任选其一，无标记走默认）：
  「【客服】」「【研发】」「【运维】」或同义 「[[客服]]」「[[研发]]」「[[运维]]」

用法:
  cd backend && .\\venv\\Scripts\\python.exe scripts/create_enterprise_scenario_agents.py
  USE_TESTCLIENT=1   # 推荐：不经 TCP

环境变量:
  PLATFORM_BASE_URL, PLATFORM_USERNAME, PLATFORM_PASSWORD
  ENTERPRISE_SCENARIO_NAMES  可选，逗号分隔，只创建/更新列出的 Agent 名称（完整名）
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any, Dict, List, Optional, Tuple

# ---------------------------------------------------------------------------
# 图构建辅助
# ---------------------------------------------------------------------------

DEFAULT_LLM_PROVIDER = os.getenv("ENTERPRISE_LLM_PROVIDER", "deepseek")
DEFAULT_LLM_MODEL = os.getenv("ENTERPRISE_LLM_MODEL", "deepseek-chat")
DEFAULT_LLM_TIMEOUT = max(30, int(os.getenv("ENTERPRISE_LLM_TIMEOUT", "180")))


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 or s == t:
            continue
        key = (s, t, e.get("sourceHandle") or "")
        if key in seen:
            continue
        seen.add(key)
        ne = dict(e)
        if not ne.get("targetHandle"):
            ne["targetHandle"] = "left"
        if not ne.get("id"):
            sh = ne.get("sourceHandle") or "r"
            ne["id"] = f"e_{s}_{t}_{sh}"
        out.append(ne)
    return out


def _llm(
    nid: str,
    pos: Tuple[int, int],
    label: str,
    prompt: str,
    *,
    temperature: float = 0.3,
    tools: Optional[List[str]] = None,
    enable_tools: bool = False,
) -> Dict[str, Any]:
    tlist = list(tools or [])
    en = bool(enable_tools and tlist)
    return {
        "id": nid,
        "type": "llm",
        "position": {"x": pos[0], "y": pos[1]},
        "data": {
            "label": label,
            "prompt": prompt,
            "provider": DEFAULT_LLM_PROVIDER,
            "model": DEFAULT_LLM_MODEL,
            "temperature": float(temperature),
            "request_timeout": DEFAULT_LLM_TIMEOUT,
            "enable_tools": en,
            "tools": tlist if en else [],
            "selected_tools": tlist if en else [],
        },
    }


def _linear_start_llm_end(
    llm_id: str,
    llm_pos: Tuple[int, int],
    llm_label: str,
    prompt: str,
    *,
    temperature: float = 0.3,
    tools: Optional[List[str]] = None,
    enable_tools: bool = False,
) -> Dict[str, Any]:
    nodes = [
        {"id": "start-1", "type": "start", "position": {"x": 80, "y": 200}, "data": {"label": "开始"}},
        _llm(llm_id, llm_pos, llm_label, prompt, temperature=temperature, tools=tools, enable_tools=enable_tools),
        {"id": "end-1", "type": "end", "position": {"x": llm_pos[0] + 240, "y": 200}, "data": {"label": "结束"}},
    ]
    edges = _sanitize_edges(
        [
            {"source": "start-1", "target": llm_id, "sourceHandle": "right", "targetHandle": "left"},
            {"source": llm_id, "target": "end-1", "sourceHandle": "right", "targetHandle": "left"},
        ]
    )
    return {"nodes": nodes, "edges": edges}


# ---------------------------------------------------------------------------
# 各场景 workflow
# ---------------------------------------------------------------------------

WF_ECOM = _linear_start_llm_end(
    "llm-1",
    (360, 200),
    "售后处理",
    """你是电商售后场景助手。处理订单状态咨询、物流、退换货政策说明；不要编造订单号与物流单号。
可调用工具辅助：HTTP（查公开接口时）、文本分析、时间。回答简洁、分条，必要时先澄清一项关键信息。""",
    temperature=0.35,
    tools=["http_request", "text_analyze", "datetime"],
    enable_tools=True,
)

WF_DEV = _linear_start_llm_end(
    "llm-1",
    (360, 200),
    "联调助手",
    """你是研发联调助手：帮助解析 API 返回、构造示例请求体、阅读本地说明文件（路径由用户提供）、简单计算。
务必以工具返回为准，勿编造响应体。优先给出可复制的 curl 或 JSON 片段。""",
    temperature=0.25,
    tools=["http_request", "json_process", "file_read", "math_calculate"],
    enable_tools=True,
)

WF_DATA = _linear_start_llm_end(
    "llm-1",
    (360, 200),
    "数据问答",
    """你是企业数据问答助手：在用户提供明确需求时，可用只读 SQL（SELECT）分析业务表；可解析 JSON、做简单统计与文本摘要。
不得执行删改；无数据源或查询失败时如实说明。回答用中文，结论在前。""",
    temperature=0.28,
    tools=["database_query", "json_process", "text_analyze", "math_calculate"],
    enable_tools=True,
)

WF_MOBILE = _linear_start_llm_end(
    "llm-1",
    (360, 200),
    "端侧运维",
    """你是移动端/设备侧运维助手：在用户需要时协助解读 adb 日志、读取用户指定的日志文件、结合系统信息排查。
若环境无 adb 或权限不足，根据工具错误如实说明。不要编造日志内容。""",
    temperature=0.3,
    tools=["adb_log", "file_read", "text_analyze", "system_info"],
    enable_tools=True,
)

WF_COMPLIANCE = _linear_start_llm_end(
    "llm-1",
    (360, 200),
    "合规纪要",
    """你是合规与对内纪要助手。输出应：条理清晰、避免绝对化承诺、敏感处提示「需法务/合规复核」；不编造内部制度条文。
不使用外部工具，基于用户给定材料与问题作答。""",
    temperature=0.2,
    tools=[],
    enable_tools=False,
)

WF_SALES = _linear_start_llm_end(
    "llm-1",
    (360, 200),
    "销售助理",
    """你是 B 端销售助理：协助整理客户痛点、撰写简短跟进话术、摘要公开产品资料（若用户提供 URL 可用 HTTP 拉取）。
保持专业、不夸大效果；数字与条款以工具或用户提供的原文为准。""",
    temperature=0.35,
    tools=["http_request", "text_analyze", "datetime"],
    enable_tools=True,
)

CODE_LANE = """
d = input_data if isinstance(input_data, dict) else {}
q = str(d.get("query", "") or "")
lane = "default"
if "【研发】" in q or "[[研发]]" in q:
    lane = "dev"
elif "【运维】" in q or "[[运维]]" in q:
    lane = "ops"
elif "【客服】" in q or "[[客服]]" in q:
    lane = "cs"
result = dict(d)
result["lane"] = lane
""".strip()


def build_switch_multilane_workflow() -> Dict[str, Any]:
    nodes: List[Dict[str, Any]] = [
        {"id": "start-1", "type": "start", "position": {"x": 60, "y": 300}, "data": {"label": "开始"}},
        {
            "id": "code-lane",
            "type": "code",
            "position": {"x": 260, "y": 300},
            "data": {"label": "解析线路", "language": "python", "code": CODE_LANE, "timeout": 15},
        },
        {
            "id": "sw-1",
            "type": "switch",
            "position": {"x": 500, "y": 300},
            "data": {
                "label": "业务线路",
                "field": "lane",
                "cases": {"cs": "br_cs", "dev": "br_dev", "ops": "br_ops"},
                "default": "br_default",
            },
        },
        _llm(
            "llm-cs",
            (760, 120),
            "客服线",
            "你是企业客服专家。用户已从多线路由进入本分支；忽略线路标记，专注解决问题。简洁、礼貌。",
            temperature=0.35,
        ),
        _llm(
            "llm-def",
            (760, 260),
            "默认线",
            "你是通用企业助手。用户未指定【客服】/【研发】/【运维】线路；先简要澄清所属场景再回答。",
            temperature=0.35,
        ),
        _llm(
            "llm-dev",
            (760, 400),
            "研发线",
            "你是研发支持专家。用户已进入研发分支；给出可执行步骤与示例，避免空泛。",
            temperature=0.28,
        ),
        _llm(
            "llm-ops",
            (760, 540),
            "运维线",
            "你是运维专家。用户已进入运维分支；优先给排查顺序与注意事项，不编造监控数据。",
            temperature=0.3,
        ),
        {
            "id": "merge-1",
            "type": "merge",
            "position": {"x": 1060, "y": 330},
            "data": {"label": "合并", "mode": "merge_all", "strategy": "object"},
        },
        {"id": "end-1", "type": "end", "position": {"x": 1300, "y": 330}, "data": {"label": "结束"}},
    ]
    edges = _sanitize_edges(
        [
            {"source": "start-1", "target": "code-lane", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "code-lane", "target": "sw-1", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "sw-1", "target": "llm-cs", "sourceHandle": "br_cs", "targetHandle": "left"},
            {"source": "sw-1", "target": "llm-dev", "sourceHandle": "br_dev", "targetHandle": "left"},
            {"source": "sw-1", "target": "llm-ops", "sourceHandle": "br_ops", "targetHandle": "left"},
            {"source": "sw-1", "target": "llm-def", "sourceHandle": "br_default", "targetHandle": "left"},
            {"source": "llm-cs", "target": "merge-1", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "llm-dev", "target": "merge-1", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "llm-ops", "target": "merge-1", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "llm-def", "target": "merge-1", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "merge-1", "target": "end-1", "sourceHandle": "right", "targetHandle": "left"},
        ]
    )
    return {"nodes": nodes, "edges": edges}


def build_hitl_approval_workflow() -> Dict[str, Any]:
    nodes: List[Dict[str, Any]] = [
        {"id": "start-1", "type": "start", "position": {"x": 80, "y": 220}, "data": {"label": "开始"}},
        {
            "id": "appr-1",
            "type": "approval",
            "position": {"x": 320, "y": 220},
            "data": {
                "label": "人工审批",
                "message": "请确认是否允许本助手根据用户 query 生成对外说明草稿（示例 HITL）。",
                "approved_branch": "approved",
                "rejected_branch": "rejected",
            },
        },
        _llm(
            "llm-ok",
            (580, 120),
            "通过后生成",
            "审批已通过。请根据用户问题写一段简短、可发送给业务方的说明草稿（中文）。",
            temperature=0.35,
        ),
        {
            "id": "end-ok",
            "type": "end",
            "position": {"x": 860, "y": 120},
            "data": {"label": "结束(通过)"},
        },
        {
            "id": "end-rej",
            "type": "end",
            "position": {"x": 580, "y": 320},
            "data": {"label": "结束(拒绝)"},
        },
    ]
    edges = _sanitize_edges(
        [
            {"source": "start-1", "target": "appr-1", "sourceHandle": "right", "targetHandle": "left"},
            {"source": "appr-1", "target": "llm-ok", "sourceHandle": "approved", "targetHandle": "left"},
            {"source": "appr-1", "target": "end-rej", "sourceHandle": "rejected", "targetHandle": "left"},
            {"source": "llm-ok", "target": "end-ok", "sourceHandle": "right", "targetHandle": "left"},
        ]
    )
    return {"nodes": nodes, "edges": edges}


ScenarioSpec = Tuple[str, str, Dict[str, Any], Dict[str, Any]]

DEFAULT_BUDGET = {"max_steps": 100, "max_llm_invocations": 8, "max_tool_calls": 40}
SWITCH_BUDGET = {"max_steps": 160, "max_llm_invocations": 12, "max_tool_calls": 48}
HITL_BUDGET = {"max_steps": 80, "max_llm_invocations": 6, "max_tool_calls": 20}


def all_scenario_specs() -> List[ScenarioSpec]:
    return [
        (
            "企业场景_电商售后",
            "电商售后：订单/物流/退换货咨询，带 HTTP·文本·时间工具。",
            WF_ECOM,
            DEFAULT_BUDGET,
        ),
        (
            "企业场景_研发联调",
            "研发联调：API/JSON/读文件/计算，适合排障与样例构造。",
            WF_DEV,
            DEFAULT_BUDGET,
        ),
        (
            "企业场景_数据分析",
            "数据分析：只读 SQL + JSON/文本/数学（依赖平台数据源）。",
            WF_DATA,
            DEFAULT_BUDGET,
        ),
        (
            "企业场景_移动运维",
            "移动运维：adb、读文件、文本与系统信息（依赖运行环境）。",
            WF_MOBILE,
            DEFAULT_BUDGET,
        ),
        (
            "企业场景_合规纪要",
            "合规纪要：无工具，偏对内记录与风险提示。",
            WF_COMPLIANCE,
            HITL_BUDGET,
        ),
        (
            "企业场景_销售助理",
            "销售助理：话术与资料摘要，带 HTTP·文本·时间。",
            WF_SALES,
            DEFAULT_BUDGET,
        ),
        (
            "企业场景_多线路由",
            "多线路由：query 含【客服】/【研发】/【运维】→ Switch → 专精 LLM → Merge。",
            build_switch_multilane_workflow(),
            SWITCH_BUDGET,
        ),
        (
            "企业场景_审批流样例",
            "审批样例：首跑 awaiting_approval；通过后 LLM 写说明，拒绝直接结束。",
            build_hitl_approval_workflow(),
            HITL_BUDGET,
        ),
    ]


def _validate_local(wf: Dict[str, Any]) -> None:
    from app.services.workflow_validator import validate_workflow

    r = validate_workflow(wf.get("nodes") or [], wf.get("edges") or [])
    if not r.get("valid"):
        errs = r.get("errors") or []
        raise ValueError("工作流校验失败: " + "; ".join(errs))


def _filter_specs(
    specs: List[ScenarioSpec], only_names: Optional[List[str]]
) -> List[ScenarioSpec]:
    if not only_names:
        return specs
    want = {n.strip() for n in only_names if n.strip()}
    return [s for s in specs if s[0] in want]


def _upsert_agent_requests(
    base: str,
    headers: Dict[str, str],
    name: str,
    description: str,
    wf: Dict[str, Any],
    budget: Dict[str, Any],
) -> Tuple[str, bool]:
    import requests

    def find_id() -> Optional[str]:
        rr = requests.get(
            f"{base}/api/v1/agents", params={"search": name, "limit": 80}, headers=headers, timeout=45
        )
        if rr.status_code != 200:
            return None
        for a in rr.json() or []:
            if a.get("name") == name:
                return a.get("id")
        return None

    existing = find_id()
    if existing:
        ur = requests.put(
            f"{base}/api/v1/agents/{existing}",
            headers=headers,
            json={"description": description, "workflow_config": wf, "budget_config": budget},
            timeout=120,
        )
        if ur.status_code != 200:
            raise RuntimeError(f"更新 {name} 失败 {ur.status_code}: {ur.text[:600]}")
        return existing, False
    cr = requests.post(
        f"{base}/api/v1/agents",
        headers=headers,
        json={"name": name, "description": description, "workflow_config": wf, "budget_config": budget},
        timeout=120,
    )
    if cr.status_code != 201:
        raise RuntimeError(f"创建 {name} 失败 {cr.status_code}: {cr.text[:600]}")
    return cr.json()["id"], True


def _upsert_agent_testclient(
    c: Any,
    headers: Dict[str, str],
    name: str,
    description: str,
    wf: Dict[str, Any],
    budget: Dict[str, Any],
) -> Tuple[str, bool]:
    existing = None
    gr = c.get("/api/v1/agents", params={"search": name, "limit": 80}, headers=headers)
    if gr.status_code == 200:
        for a in gr.json() or []:
            if a.get("name") == name:
                existing = a.get("id")
                break
    if existing:
        ur = c.put(
            f"/api/v1/agents/{existing}",
            headers=headers,
            json={"description": description, "workflow_config": wf, "budget_config": budget},
        )
        if ur.status_code != 200:
            raise RuntimeError(f"更新 {name} 失败 {ur.status_code}: {ur.text[:600]}")
        return existing, False
    cr = c.post(
        "/api/v1/agents",
        headers=headers,
        json={"name": name, "description": description, "workflow_config": wf, "budget_config": budget},
    )
    if cr.status_code != 201:
        raise RuntimeError(f"创建 {name} 失败 {cr.status_code}: {cr.text[:600]}")
    return cr.json()["id"], True


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(description="批量创建企业场景 Agent")
    p.add_argument(
        "--names",
        nargs="*",
        help="仅处理这些完整 Agent 名称（可替代环境变量 ENTERPRISE_SCENARIO_NAMES）",
    )
    args = p.parse_args(argv)

    env_filter = os.getenv("ENTERPRISE_SCENARIO_NAMES", "").strip()
    only_list: Optional[List[str]] = None
    if args.names:
        only_list = list(args.names)
    elif env_filter:
        only_list = [x.strip() for x in env_filter.split(",") if x.strip()]

    specs = _filter_specs(all_scenario_specs(), only_list)

    results: List[Dict[str, Any]] = []
    errors: List[str] = []

    use_tc = os.getenv("USE_TESTCLIENT", "").strip().lower() in ("1", "true", "yes")

    if use_tc:
        from fastapi.testclient import TestClient

        from app.main import app

        c = TestClient(app)
        lr = c.post(
            "/api/v1/auth/login",
            data={
                "username": os.getenv("PLATFORM_USERNAME", "admin"),
                "password": os.getenv("PLATFORM_PASSWORD", "123456"),
            },
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )
        if lr.status_code != 200:
            print("login:", lr.status_code, lr.text[:400], file=sys.stderr)
            return 1
        token = lr.json().get("access_token")
        headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

        for name, desc, wf, budget in specs:
            try:
                _validate_local(wf)
                aid, created = _upsert_agent_testclient(c, headers, name, desc, wf, budget)
                results.append({"name": name, "id": aid, "created": created})
                tag = "创建" if created else "更新"
                print(f"[{tag}] {name}  {aid}")
            except Exception as e:
                errors.append(f"{name}: {e}")
                print(f"[失败] {name}: {e}", file=sys.stderr)
    else:
        import requests

        base = os.getenv("PLATFORM_BASE_URL", "http://127.0.0.1:8037").rstrip("/")
        lr = requests.post(
            f"{base}/api/v1/auth/login",
            data={
                "username": os.getenv("PLATFORM_USERNAME", "admin"),
                "password": os.getenv("PLATFORM_PASSWORD", "123456"),
            },
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            timeout=20,
        )
        if lr.status_code != 200:
            print("登录失败:", lr.status_code, lr.text[:400], file=sys.stderr)
            return 1
        token = lr.json().get("access_token")
        if not token:
            print("无 access_token", file=sys.stderr)
            return 1
        headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

        for name, desc, wf, budget in specs:
            try:
                _validate_local(wf)
                aid, created = _upsert_agent_requests(base, headers, name, desc, wf, budget)
                results.append({"name": name, "id": aid, "created": created})
                tag = "创建" if created else "更新"
                print(f"[{tag}] {name}  {aid}")
            except Exception as e:
                errors.append(f"{name}: {e}")
                print(f"[失败] {name}: {e}", file=sys.stderr)

    summary = {"ok": len(results), "fail": len(errors), "agents": results, "errors": errors}
    print(json.dumps(summary, ensure_ascii=False))
    return 0 if not errors else 2


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