"""
nl2agent.py — 一句话搭建智能体

Natural Language → YAML Config → Lambda Term → Execute

用法:
    python nl2agent.py "帮我搭一个深度调研助手，能搜索、写代码、做计算，最多20步，带记忆"
    python nl2agent.py --interactive   # 交互模式
"""
from __future__ import annotations

import os
import sys
import json
import yaml
import time
import tempfile
from typing import Optional

# ============================================================
# 1. NL → YAML: 用 LLM 将自然语言描述编译为 agent 配置
# ============================================================

SYSTEM_PROMPT = '''你是一个 Agent 配置编译器。用户用自然语言描述需要什么样的智能体，你输出对应的 YAML 配置。

你必须严格输出合法 YAML，不要输出任何其他内容（无解释、无 markdown 代码块标记）。

可用的 agent type:
- simple: 单步对话 agent
- react: 多步推理 agent (带工具调用, Y 组合子循环)
- chain: 流水线 agent (多步顺序执行)
- router: 路由 agent (根据输入分类到不同子 agent)
- parallel: 并行 agent (多个 agent 同时执行后合并)

可用的 MCP 工具 (onlineTool 下 example-mcp-server 服务器):
- everything_get_sum: 搜索工具，可搜索互联网信息
- chat_improve_prompt: 优化 prompt 的工具

本地工具 (localTools):
- terminate: 结束任务（react 类型必须包含）

模型选择:
- dashscope 的 qwen3-max-2026-01-23（默认）
- anthropic 的 claude-sonnet-4-20250514

配置结构参考:
```yaml
agentId: string
name: string
description: string
type: react|chain|router|parallel|simple
systemPrompt: |
  系统提示词...
model:
  provider: dashscope
  name: qwen3-max-2026-01-23
  temperature: 0.7
  maxTokens: 4096
react:  # type=react 时
  maxSteps: 10
  observationEnabled: true
  toolTimeout: 30
chain:  # type=chain 时
  steps:
    - name: step_name
      prompt: "步骤提示词"
router:  # type=router 时
  classifier:
    prompt: "分类提示词"
    categories: [cat1, cat2]
  routes:
    cat1: {type: simple, systemPrompt: "..."}
    cat2: {type: simple, systemPrompt: "..."}
  default: {type: simple, systemPrompt: "..."}
parallel:  # type=parallel 时
  agents:
    - {name: agent1, systemPrompt: "..."}
    - {name: agent2, systemPrompt: "..."}
  merge: custom
  mergePrompt: "合并提示词"
memory:
  enabled: true
  strategy: local
  size: 20
  ttl: 7200
guard:
  validator: "len(x) > 100"
  retry: 1
mcp:
  onlineTool:
    example-mcp-server:
      - everything_get_sum
  localTools:
    - terminate
  policy:
    mode: auto
app:
  mcp:
    custom:
      nodes:
        example-mcp-server:
          url: https://your-mcp-endpoint.example.com
          endpoint: /mcp/airouting
          headers:
            Authorization: "${MCP_AUTH_TOKEN}"
```

规则:
1. react 类型必须在 localTools 中包含 terminate
2. systemPrompt 要详细、专业、有针对性
3. 根据用户描述的复杂度选择合适的 type
4. 如果用户提到"搜索"、"查资料"，用 everything_get_sum 工具
5. 如果用户提到多个独立视角/角度，考虑用 parallel
6. 如果用户提到步骤/流程，考虑用 chain
7. 如果用户提到分类/路由/不同情况，考虑用 router
8. app.mcp.custom.nodes 的配置固定不变（如上面的参考）
9. 只输出 YAML，不要任何其他文字'''


def nl_to_yaml(description: str, model: str = "qwen3-max-2026-01-23") -> str:
    """
    用 LLM 将自然语言描述转为 YAML 配置。
    Lambda 语义: nl_to_yaml = λ(description). LLM_{compiler}(description)
    """
    provider = _detect_provider(model)

    if provider == "dashscope":
        return _call_dashscope(model, SYSTEM_PROMPT, description)
    elif provider == "anthropic":
        return _call_anthropic(model, SYSTEM_PROMPT, description)
    else:
        return _call_dashscope("qwen3-max-2026-01-23", SYSTEM_PROMPT, description)


def _detect_provider(model: str) -> str:
    if "qwen" in model.lower() or "dashscope" in model.lower():
        return "dashscope"
    elif "claude" in model.lower() or "anthropic" in model.lower():
        return "anthropic"
    return "dashscope"


def _call_dashscope(model, system, user) -> str:
    import urllib.request
    api_key = os.environ.get("DASHSCOPE_API_KEY", "")
    if not api_key:
        raise RuntimeError("请设置 DASHSCOPE_API_KEY 环境变量")

    url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
    body = json.dumps({
        "model": model,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "temperature": 0.3,
        "max_tokens": 4096,
    }).encode("utf-8")

    req = urllib.request.Request(url, data=body, headers={
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    })
    with urllib.request.urlopen(req, timeout=60) as resp:
        data = json.loads(resp.read().decode("utf-8"))
    return data["choices"][0]["message"]["content"].strip()


def _call_anthropic(model, system, user) -> str:
    import anthropic
    client = anthropic.Anthropic()
    resp = client.messages.create(
        model=model, max_tokens=4096, temperature=0.3,
        system=system,
        messages=[{"role": "user", "content": user}],
    )
    return resp.content[0].text.strip()


# ============================================================
# 2. YAML → Lambda Term → Execute
# ============================================================

def build_and_run(yaml_str: str, user_input: str, verbose: bool = True) -> str:
    """
    从 YAML 字符串编译 Agent 并执行。

    完整流程:
        YAML string → parse → Term → term(input) → result
                       |        |         |
                    解析      编译     β-规约
    """
    from lambdagent.fromconfig import from_config, lint_config, format_lint
    from lambdagent.fromconfig import to_lambda_expr, describe_config

    # 写入临时文件
    with tempfile.NamedTemporaryFile(mode='w', suffix='.yml', delete=False, encoding='utf-8') as f:
        f.write(yaml_str)
        tmp_path = f.name

    try:
        # Lint
        if verbose:
            results = lint_config(tmp_path)
            print(format_lint(results, "generated-agent.yml"))
            print()

        # Lambda 结构
        if verbose:
            print(describe_config(tmp_path))
            print()
            print("Lambda expression:")
            print(to_lambda_expr(tmp_path))
            print()

        # 编译
        if verbose:
            print("=" * 60)
            print("  Compiling YAML → Lambda Term...")
        agent = from_config(tmp_path)
        if verbose:
            print(f"  Compiled: {agent}")
            print("=" * 60)
            print()

        # 执行
        if verbose:
            print(f"  Executing: agent(\"{user_input[:50]}...\")")
            print("  " + "-" * 56)

        t0 = time.time()
        result = agent(user_input)
        duration = time.time() - t0

        if verbose:
            print(f"\n  Done in {duration:.1f}s")
            print("=" * 60)

        return result

    finally:
        os.unlink(tmp_path)


# ============================================================
# 3. 完整流程: 一句话 → Agent
# ============================================================

def one_sentence_to_agent(description: str, user_input: Optional[str] = None,
                          model: str = "qwen3-max-2026-01-23", verbose: bool = True):
    """
    一句话搭建并运行 Agent。

    Lambda 语义:
        one_sentence_to_agent = λ(desc, input).
            let yaml  = LLM_{compiler}(desc) in        ← 意图解析
            let term  = from_config(yaml) in            ← 编译
            term(input)                                  ← β-规约
    """
    print("=" * 60)
    print("  lambdagent — 一句话搭建智能体")
    print("=" * 60)
    print()
    print(f"  描述: \"{description}\"")
    print()

    # Step 1: NL → YAML
    print("  [Step 1] 解析意图，生成配置...")
    t0 = time.time()
    yaml_str = nl_to_yaml(description, model)

    # 清理: 去掉可能的 markdown 代码块标记
    yaml_str = yaml_str.strip()
    if yaml_str.startswith("```"):
        lines = yaml_str.split("\n")
        # Remove first and last lines if they are code fences
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        yaml_str = "\n".join(lines)

    gen_time = time.time() - t0
    print(f"  Done ({gen_time:.1f}s)")
    print()

    # 显示生成的 YAML
    print("  [Generated YAML Config]")
    print("  " + "-" * 56)
    for line in yaml_str.split("\n"):
        print(f"  {line}")
    print("  " + "-" * 56)
    print()

    # 验证 YAML 合法性
    try:
        cfg = yaml.safe_load(yaml_str)
        if not isinstance(cfg, dict):
            print("  [ERROR] Generated YAML is not a valid config dict")
            return yaml_str
    except yaml.YAMLError as e:
        print(f"  [ERROR] Invalid YAML: {e}")
        return yaml_str

    # Step 2: 保存 YAML
    config_path = "generated_agent.yml"
    with open(config_path, "w", encoding="utf-8") as f:
        f.write(yaml_str)
    print(f"  Config saved to: {config_path}")
    print()

    # Step 3: Lint + Lambda 结构
    from lambdagent.fromconfig import lint_config, format_lint, describe_config, to_lambda_expr

    results = lint_config(cfg)
    print(format_lint(results, "generated_agent.yml"))
    print()

    print(describe_config(cfg))
    print()

    print("  [Lambda Expression]")
    print(f"  {to_lambda_expr(cfg)}")
    print()

    # Step 4: 编译
    from lambdagent.fromconfig import from_config
    print("  [Step 2] Compiling YAML → Lambda Term...")
    agent = from_config(config_path)
    print(f"  Compiled: {agent}")
    print()

    # Step 5: 如果有输入就执行
    if user_input:
        print(f"  [Step 3] Executing agent(\"{user_input[:60]}\")")
        print("  " + "-" * 56)
        t0 = time.time()
        result = agent(user_input)
        duration = time.time() - t0
        print()
        print(f"  [Result] ({duration:.1f}s)")
        print("  " + "-" * 56)
        print(result)
        print("  " + "-" * 56)
        return result
    else:
        print("  Agent compiled successfully! Use agent(input) to execute.")
        print()
        print("  Example:")
        print(f'    agent("{description[:30]}相关的问题")')
        return agent


# ============================================================
# 4. 交互模式
# ============================================================

def interactive_mode(model: str = "qwen3-max-2026-01-23"):
    """交互式一句话搭建 Agent。"""
    print("=" * 60)
    print("  lambdagent 交互模式")
    print("  输入自然语言描述来搭建智能体")
    print("  输入 :quit 退出")
    print("=" * 60)
    print()

    while True:
        try:
            desc = input("描述你需要的 Agent:\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye!")
            break

        if not desc or desc == ":quit":
            print("Bye!")
            break

        # 询问是否有具体任务
        try:
            task = input("要执行什么任务? (直接回车跳过):\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye!")
            break

        print()
        one_sentence_to_agent(desc, task or None, model)
        print()


# ============================================================
# 5. 预置 Demo
# ============================================================

DEMO_DESCRIPTIONS = {
    "research": (
        "帮我搭一个深度调研助手，能用搜索工具查资料，能分析整理信息，"
        "如果信息不够就继续搜索，最多研究20步，带记忆功能，"
        "最后输出一份完整的、超过500字的调研报告",
        "请调研：大语言模型在软件工程领域的最新应用进展，包括代码生成、代码审查、自动化测试等方面"
    ),
    "customer_service": (
        "搭一个智能客服系统，能自动识别用户是咨询技术问题、账号问题还是投诉建议，"
        "分别路由给不同的专家处理，技术专家能调用搜索工具查文档",
        "我的账号登录不了了，显示密码错误，但我确认密码是对的"
    ),
    "code_review": (
        "做一个代码审查流水线，第一步检查安全漏洞，第二步检查代码风格和可读性，"
        "第三步检查性能问题，第四步生成综合审查报告，每一步的输出要超过100字",
        "def login(user, pwd):\n  q = f\"SELECT * FROM users WHERE name='{user}' AND pass='{pwd}'\"\n  return db.execute(q)"
    ),
    "translator": (
        "做一个多语言翻译系统，把输入同时翻译成英文、日文、韩文三个版本并行执行，"
        "然后由一个总结专家对比三个翻译版本的质量并给出最终推荐",
        "人工智能正在深刻改变软件开发的方式，从代码自动生成到智能调试，开发者的工作效率得到了显著提升。"
    ),
}


def run_demo(name: str = "research", model: str = "qwen3-max-2026-01-23"):
    """运行预置 Demo。"""
    if name not in DEMO_DESCRIPTIONS:
        print(f"Available demos: {list(DEMO_DESCRIPTIONS.keys())}")
        return

    desc, task = DEMO_DESCRIPTIONS[name]
    return one_sentence_to_agent(desc, task, model)


# ============================================================
# Main
# ============================================================

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="一句话搭建智能体")
    parser.add_argument("description", nargs="?", default=None,
                        help="Agent 的自然语言描述")
    parser.add_argument("--task", "-t", default=None,
                        help="要执行的任务")
    parser.add_argument("--interactive", "-i", action="store_true",
                        help="交互模式")
    parser.add_argument("--demo", "-d", default=None,
                        choices=list(DEMO_DESCRIPTIONS.keys()),
                        help="运行预置 Demo")
    parser.add_argument("--model", "-m", default="qwen3-max-2026-01-23",
                        help="用于生成配置的模型")
    parser.add_argument("--generate-only", "-g", action="store_true",
                        help="只生成 YAML，不执行")
    args = parser.parse_args()

    if args.interactive:
        interactive_mode(args.model)
    elif args.demo:
        run_demo(args.demo, args.model)
    elif args.description:
        task = args.task if not args.generate_only else None
        one_sentence_to_agent(args.description, task, args.model)
    else:
        parser.print_help()
        print("\n示例:")
        print('  python nl2agent.py "搭一个能搜索和分析的调研助手" -t "调研大模型最新进展"')
        print('  python nl2agent.py --demo research')
        print('  python nl2agent.py --interactive')
        print(f'\n预置 Demo: {list(DEMO_DESCRIPTIONS.keys())}')