#!/usr/bin/env python3
"""最简轨迹重放（pause/resume，保留沙箱，打印每次 pause 的快照文件）。

====================================================================
使用说明
====================================================================

功能:
    基于一条 .replay.json 轨迹，在 AgentENV 上创建 1 个沙箱并回放:

        create -> pause(init)
        每条 action: sleep(delay_time) -> resume -> 执行命令 -> pause
                     -> 打印本次 pause 新增的快照文件
        结束: 不删除沙箱（保持 paused），打印 sandbox_id 和快照目录

    每次 pause 的快照文件由 AgentENV 写入宿主机本地目录:

        /var/lib/aenv/persisted-sandboxes/artifacts/<sandbox_id>/<generation>/

前提:
    1. 在 AgentENV 宿主机上运行本脚本（快照文件在服务端本地磁盘）；
    2. AgentENV 中已有 ready 状态的 template（如 django-money-task-v2）；
    3. 已安装 e2b SDK（本仓库: uv pip install -e .）；
    4. 已设置连接环境变量，例如单机无鉴权部署:

        export E2B_API_URL='http://127.0.0.1:8000'
        export E2B_SANDBOX_URL="$E2B_API_URL"
        export E2B_API_KEY='e2b_0000000000000000000000000000000000000000'
        export E2B_ACCESS_TOKEN='dummy'

用法:
    python agent/replay_pause_min.py <trajectory.replay.json> [template] [workdir]

参数:
    trajectory.replay.json   轨迹文件（必填）
    template                 AgentENV 模板名（默认 django-money-task-v2）
    workdir                  沙箱内执行目录（默认 /testbed）

示例:
    # 1) 最简：只用默认值
    python agent/replay_pause_min.py \
        delay_time_trajectories/django-money__django-money.835c1ab8.combine_file__i1ay87eb.replay.json

    # 2) 显式指定 template 和 workdir
    python agent/replay_pause_min.py \
        delay_time_trajectories/django-money__django-money.835c1ab8.combine_file__i1ay87eb.replay.json \
        django-money-task-v2 /testbed

    # 3) 跑完后确认沙箱仍在（应为 paused），快照文件在磁盘上
    aenv list --output table
    ls -l /var/lib/aenv/persisted-sandboxes/artifacts/<sandbox_id>/

    # 4) 确认无误后手动清理沙箱
    aenv delete <sandbox_id>

输出说明:
    [pause <tag>] 行后面逐行列出本次 pause 新增的快照文件:
        <字节数>  <文件绝对路径>
    tag 为 init 或 step-0001、step-0002 ...，对应轨迹中的第 N 条 action。

注意:
    - 这是最简版本：没有重试、没有 instance patch 初始化、没有结果落盘；
      需要更完整的版本请用同目录的 replay_pause_keep.py。
    - 如想要干净的快照基线，先在确认无其他任务后清空
      /var/lib/aenv/persisted-sandboxes 下的旧数据。
"""

import asyncio
import json
import os
import shlex
import sys
import time
from pathlib import Path

ART = Path("/var/lib/aenv/persisted-sandboxes/artifacts")
STOP = {"submit", "finish", "done"}


def gens(sid: str) -> list[str]:
    base = ART / sid
    return sorted(e.name for e in os.scandir(base) if e.is_dir(follow_symlinks=False)) if base.is_dir() else []


def gen_files(sid: str, gen: str) -> list[str]:
    base = ART / sid / gen
    return [str(p) for p in sorted(base.rglob("*")) if p.is_file()]


def load_steps(path: str) -> list[tuple[str, float]]:
    data = json.loads(Path(path).read_text(encoding="utf-8"))
    steps = []
    for raw in data.get("trajectory", []):
        action = (raw.get("action") or "").strip() if isinstance(raw, dict) else ""
        if not action:
            continue
        if action.lower() in STOP:
            break
        delay = raw.get("delay_time", 0.0)
        steps.append((action, float(delay) if isinstance(delay, (int, float)) and delay > 0 else 0.0))
    if not steps:
        sys.exit(f"no executable actions in {path}")
    return steps


async def main() -> None:
    traj = sys.argv[1]
    template = sys.argv[2] if len(sys.argv) > 2 else "django-money-task-v2"
    workdir = sys.argv[3] if len(sys.argv) > 3 else "/testbed"
    steps = load_steps(traj)

    from e2b import AsyncSandbox

    sb = await AsyncSandbox.create(template=template, timeout=86400, request_timeout=120)
    sid = sb.sandbox_id
    print(f"sandbox: {sid}  steps: {len(steps)}", flush=True)
    seen = set(gens(sid))

    async def pause_collect(tag: str) -> None:
        pause = getattr(sb, "pause", None) or sb.beta_pause
        t0 = time.perf_counter()
        await pause(request_timeout=120)
        new = [g for g in gens(sid) if g not in seen]
        seen.update(new)
        files = [f for g in new for f in gen_files(sid, g)]
        print(f"[pause {tag}] {time.perf_counter() - t0:.2f}s, new files:", flush=True)
        for f in files:
            print(f"  {os.path.getsize(f):>12}  {f}", flush=True)

    await pause_collect("init")
    for i, (action, delay) in enumerate(steps, 1):
        if delay:
            await asyncio.sleep(delay)
        sb = await sb.connect(timeout=86400, request_timeout=120)  # resume
        cmd = f"bash -lc {shlex.quote(f'cd {shlex.quote(workdir)}; {action}')}"
        handle = await sb.commands.run(cmd=cmd, background=True, timeout=300, user="root")
        try:
            r = await handle.wait()
        except Exception as exc:  # 非零退出码会以异常形式抛出
            r = exc
        print(f"[{i}/{len(steps)}] exit={getattr(r, 'exit_code', '?')}  {action[:80]}", flush=True)
        await pause_collect(f"step-{i:04d}")

    # 不 kill：沙箱保持 paused，快照文件留在磁盘上
    print(f"\nkept sandbox (paused): {sid}", flush=True)
    print(f"snapshot files: {ART / sid}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())