#!/usr/bin/env python3
"""agent_loop_demo.py — 用确定性模拟演示 Claude Code 的 agentic loop。

不需要 API key：用一个"假装的模型"( scripted policy ) 代替真实 Claude，
复现官方文档描述的三阶段循环：gather context -> take action -> verify results。
（官方定义来源: https://code.claude.com/docs/en/how-claude-code-works , 访问于 2026-08-30）

运行: python3 agent_loop_demo.py [任务文件]
"""

from __future__ import annotations

import sys
from dataclasses import dataclass, field


# ---------- 模拟项目环境 ----------

@dataclass
class FakeRepo:
    files: dict[str, str] = field(default_factory=dict)
    test_results: dict[str, bool] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.files = {
            "src/auth.py": (
                "def login(user, pwd):\n"
                "    token = make_t…r)\n"
                "    return token\n"
            ),
            "src/tokens.py": (
                "def make_token(user):\n"
                "    # BUG: token 过期时间被写死为 1 秒，刷新后立刻 401\n"
                "    return {'exp': 1}\n",
            ),
            "tests/test_auth.py": (
                "def test_token_expiry():\n"
                "    t = make_token('u')\n"
                "    assert t['exp'] > 60, 'token must live > 60s'\n",
            ),
        }
        self.test_results = {"tests/test_auth.py": False}


# ---------- 模拟工具（对应 Claude Code 内置工具类别） ----------

class Tools:
    """每类工具都对应 Claude Code 的一个内置类别（File ops / Search / Execution）。"""

    def __init__(self, repo: FakeRepo):
        self.repo = repo
        self.calls: list[str] = []

    def read(self, path: str) -> str:
        self.calls.append(f"Read({path})")
        return self.repo.files.get(path, f"<FileNotFoundError: {path}>")

    def grep(self, pattern: str) -> list[str]:
        self.calls.append(f"Grep({pattern})")
        hits = [p for p, c in self.repo.files.items() if pattern in c]
        return hits

    def edit(self, path: str, old: str, new: str) -> str:
        self.calls.append(f"Edit({path})")
        if path not in self.repo.files or old not in self.repo.files[path]:
            return "error: old string not found"
        self.repo.files[path] = self.repo.files[path].replace(old, new)
        return "ok"

    def bash(self, cmd: str) -> str:
        self.calls.append(f"Bash({cmd})")
        if cmd == "pytest -q":
            ok = self.repo.files["src/tokens.py"].count("exp': 3600") > 0
            self.repo.test_results["tests/test_auth.py"] = ok
            return "1 passed" if ok else "1 failed: token must live > 60s"
        return f"ran: {cmd}"


# ---------- 模拟 harness：驱动 agentic loop ----------

def run_agent(task: str, tools: Tools, max_steps: int = 12) -> None:
    """gather context -> take action -> verify results 循环。

    真实 Claude Code 中"下一步做什么"由模型决定；这里用一个小策略表模拟，
    重点展示 harness 的职责：执行工具、把结果喂回、循环直到验证通过。
    """
    context: list[str] = [f"task: {task}"]
    for step in range(1, max_steps + 1):
        action = scripted_policy(context, tools)

        if action[0] == "stop":
            print(f"\n[stop] {action[1]}")
            break

        kind, arg = action
        print(f"step {step:>2} | {kind:<10} {arg}")

        # 权限层：真实 Claude Code 在这里插入 permission 检查（默认模式会询问）
        if kind in ("edit", "bash"):
            print(f"          ↳ [permission] Manual 模式下这里会弹批准提示（本 demo 自动允许）")

        result = dispatch(kind, arg, tools)
        context.append(f"{kind}({arg}) -> {result}")
        print(f"          ↳ {result[:90]}")
    else:
        print("[max_steps] 达到步数上限仍未完成（真实系统中应触发 /compact 或升级人工）")


def scripted_policy(context: list[str], tools: Tools):
    """极简'模型'：根据上下文决定下一个动作。"""
    joined = " | ".join(context)
    if "task:" in context[0] and "grep" not in joined:
        return ("grep", "make_token")                     # gather context
    if "tests/test_auth.py" in tools.repo.files and "read tokens.py" not in joined and "exp" in joined:
        if "Read(src/tokens.py)" not in tools.calls:
            return ("read", "src/tokens.py")              # gather context（深挖）
    if "Edit(src/tokens.py)" not in tools.calls:
        return ("edit", "src/tokens.py :: {'exp': 1} -> {'exp': 3600}")  # take action
    if "1 passed" not in joined:
        return ("bash", "pytest -q")                      # verify results
    return ("stop", "测试通过，任务完成（verify 通过 -> 交付）")


def dispatch(kind: str, arg, tools: Tools) -> str:
    if kind == "grep":
        return f"hits: {tools.grep(arg)}"
    if kind == "read":
        return tools.read(arg)
    if kind == "edit":
        path, rest = arg.split(" :: ")
        old, new = rest.split(" -> ")
        return tools.edit(path, old, new)
    if kind == "bash":
        return tools.bash(arg)
    raise ValueError(kind)


def main() -> None:
    task = sys.argv[1] if len(sys.argv) > 1 else "修复 token 刷新后 401 的 bug 并跑通测试"
    repo = FakeRepo()
    tools = Tools(repo)
    print(f"=== agentic loop demo ===\ntask: {task}\n")
    run_agent(task, tools)
    print(f"\n工具调用轨迹（transcript）: {tools.calls}")
    print("把探索类调用换成 subagent，可让主上下文只保留结论——见模块 README 9.10。")


if __name__ == "__main__":
    main()