import asyncio, os, sys, json, hashlib
from pathlib import Path
from claude_agent_sdk import AssistantMessage, ToolUseBlock, ClaudeSDKClient, ClaudeAgentOptions, UserMessage, ResultMessage

W = Path(sys.argv[1]).resolve()
def snap():
    out = {}
    for p in sorted(W.rglob('*')):
        r = str(p.relative_to(W))
        if r.startswith('.claude'): continue
        if p.is_symlink(): out[r] = 'symlink->' + os.readlink(p)
        elif p.is_file(): out[r] = p.read_text()
        elif p.is_dir(): out[r] = '<dir>'
    return out

def setup():
    W.mkdir(parents=True, exist_ok=True)
    (W/'native.txt').write_text('orig\n')
    (W/'shellonly.txt').write_text('orig\n')
    (W/'mixed.txt').write_text('orig\n')
    (W/'gone.txt').write_text('orig\n')
    (W/'sub.txt').write_text('orig\n')
    (W/'target.txt').write_text('orig\n')
    os.symlink('target.txt', W/'link.txt')
    (W/'untracked-note.txt').write_text('orig\n')

PROMPT = """Do these steps exactly, using the named tool for each. Do not read files first. Reply 'done' at the end.
1. Edit tool: replace the whole content of native.txt with 'EDITED'.
2. Bash tool: run `echo BASH > shellonly.txt`.
3. Edit tool: replace the content of mixed.txt with 'EDITED'. Then Bash tool: run `echo BASH-AFTER >> mixed.txt`.
4. Write tool: create new-by-write.txt containing 'NEW'. Bash tool: `echo NEW > new-by-bash.txt`.
5. Bash tool: `rm gone.txt`.
6. Agent tool (a general-purpose subagent): ask it to use Edit/Write to replace the content of sub.txt with 'SUBAGENT'.
7. Read tool on link.txt, then Edit tool: replace the content of link.txt with 'VIA-LINK' (it is a symlink to target.txt).
8. Bash tool: `echo BASH > untracked-note.txt`."""

async def main():
    setup()
    before = snap()
    opts = ClaudeAgentOptions(cwd=str(W), enable_file_checkpointing=True, permission_mode='bypassPermissions',
        setting_sources=[], extra_args={'replay-user-messages': None, 'strict-mcp-config': None}, max_turns=30)
    ids = []
    async with ClaudeSDKClient(opts) as c:
        await c.query(PROMPT)
        async for m in c.receive_response():
            if isinstance(m, UserMessage) and m.uuid: ids.append(m.uuid)
            if isinstance(m, AssistantMessage):
                for b in m.content:
                    if isinstance(b, ToolUseBlock): print('TOOL', b.name, json.dumps(b.input)[:110])
        after = snap()
        print('checkpoint ids', len(ids))
        (W/'native.txt').write_text('HUMAN-EDIT\n')
        res = await c.rewind_files(ids[0])
        print('rewind result', res)
    rew = snap()
    rows = []
    for k in sorted(set(before)|set(after)|set(rew)):
        rows.append((k, before.get(k), after.get(k), rew.get(k)))
    for k,b,a,r in rows:
        print(f'{k:22} before={b!r:22} after={a!r:22} rewound={r!r:22} {"RESTORED" if r==b else "NOT"}')
asyncio.run(main())