#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 Claude Code 会话记录里恢复被误删的 P1 源码。

## 背景

`tools/p2_import_legacy.py` 在拷贝老代码前会 `shutil.rmtree(目标目录)` 以保证
可重复执行。但 `parking-module-system` 与 `parking-module-api` 这两个模块
**既有 P1 手写的新代码、又要接收老代码**，于是 P1 写的 84 个类被连带删掉了
（只剩 target/classes 里的 .class）。

## 恢复来源

P1 的代码是在同一个 Claude Code 会话里用 Write/Edit 工具写的，
其完整内容按顺序记录在会话 jsonl 里。本脚本按记录顺序重放：
`Write` 置内容、`Edit` 做替换，得到每个文件在被删前的最终内容。

另有 4 个 `parking-module-api` 的文件是当时用 `sed` 从老 `zteits-api`
复制+改包名生成的，内容可确定性重算，不在 jsonl 里，由
`regen_api_apibase.sh` 负责。

输出先落到 `tools/recover/out/`（不直接写源码树），核对无误后再拷回。

用法：python3 -I tools/recover/restore_p1.py <会话.jsonl> [--apply]
"""

import argparse
import collections
import json
import os
import sys

CLOUD = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
STAGE = os.path.join(CLOUD, "tools/recover/out")

# 只恢复这两个模块下 P1 手写的新代码（老代码的包名是 com.zteits.{park,order,...}
# 与 com.clouds.common，不会命中这里）
PREFIXES = (
    "parking-module-system/src/main/java/com/zteits/parking",
    "parking-module-api/src/main/java/com/zteits/parking",
)


def collect(jsonl):
    """按出现顺序收集目标文件的操作。

    ⚠️ 必须剔除**失败**的操作。会话里 Edit 报错（old_string 对不上）后助手会重试，
    若把失败的也重放一遍，重放内容就会和原始历史分叉。失败与否只能看
    `tool_result.is_error`，所以要先扫一遍拿到失败集合，再过滤。
    """
    ops = collections.OrderedDict()
    lines = []
    with open(jsonl, encoding="utf-8", errors="replace") as f:
        for line in f:
            if '"tool_use"' in line or '"tool_result"' in line:
                lines.append(line)

    failed = set()
    for line in lines:
        if '"tool_result"' not in line:
            continue
        try:
            rec = json.loads(line)
        except Exception:
            continue
        content = (rec.get("message") or {}).get("content")
        if not isinstance(content, list):
            continue
        for c in content:
            if isinstance(c, dict) and c.get("type") == "tool_result" and c.get("is_error"):
                if c.get("tool_use_id"):
                    failed.add(c["tool_use_id"])

    skipped_reads = skipped_failed = 0
    for line in lines:
        if '"tool_use"' not in line:
            continue
        try:
            rec = json.loads(line)
        except Exception:
            continue
        content = (rec.get("message") or {}).get("content")
        if not isinstance(content, list):
            continue
        for c in content:
            if not isinstance(c, dict) or c.get("type") != "tool_use":
                continue
            name = c.get("name")
            inp = c.get("input") or {}
            p = inp.get("file_path")
            if not isinstance(p, str):
                continue
            if not any(pp in p for pp in PREFIXES):
                continue
            if name == "Read":
                skipped_reads += 1
                continue
            if c.get("id") in failed:
                skipped_failed += 1
                continue
            if name not in ("Write", "Edit"):
                print(f"  !! 未处理的操作类型 {name} on {p}")
                continue
            ops.setdefault(p, []).append((name, inp))
    return ops, skipped_reads, skipped_failed


def replay(path, actions):
    txt = None
    for i, (name, inp) in enumerate(actions):
        if name == "Write":
            txt = inp.get("content", "")
        else:  # Edit
            old, new = inp.get("old_string"), inp.get("new_string")
            if txt is None:
                raise SystemExit(f"!! {path}: 第 {i} 个操作是 Edit 但没有前置 Write")
            if old is None or new is None:
                raise SystemExit(f"!! {path}: Edit 缺少 old_string/new_string")
            cnt = txt.count(old)
            if cnt == 0:
                raise SystemExit(f"!! {path}: Edit #{i} 的 old_string 在重放内容里找不到")
            if inp.get("replace_all"):
                txt = txt.replace(old, new)
            else:
                if cnt > 1:
                    print(f"  ! {path}: Edit #{i} 匹配 {cnt} 处，按语义只替换第一处")
                txt = txt.replace(old, new, 1)
    return txt


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("jsonl")
    ap.add_argument("--apply", action="store_true",
                    help="恢复后直接写回源码树（默认只落 out/ 暂存区）")
    args = ap.parse_args()

    ops, skipped_reads, skipped_failed = collect(args.jsonl)
    print(
        f"命中文件 {len(ops)} 个"
        f"（忽略 Read {skipped_reads} 次 / 失败操作 {skipped_failed} 次）"
    )

    failures = []
    written = 0
    for p, actions in ops.items():
        rel = p.split("parking-cloud/", 1)[-1]
        try:
            txt = replay(rel, actions)
        except SystemExit as e:
            failures.append(str(e))
            continue
        if txt is None:
            failures.append(f"!! {rel}: 没有 Write 内容")
            continue
        dst = os.path.join(STAGE, rel)
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        with open(dst, "w", encoding="utf-8") as f:
            f.write(txt)
        written += 1
        if args.apply:
            real = os.path.join(CLOUD, rel)
            os.makedirs(os.path.dirname(real), exist_ok=True)
            with open(real, "w", encoding="utf-8") as f:
                f.write(txt)

    print(f"恢复 {written} 个文件 -> {os.path.relpath(STAGE, CLOUD)}/")
    if args.apply:
        print("（--apply：已同时写回源码树）")
    for f in failures:
        print(f)


if __name__ == "__main__":
    main()