编辑块选择与产物恢复

"""反事实编辑信用(counterfactual edit credit)—— 纯逻辑 + 反事实 verifier 的 harbor 封装。

思想(0909 与 gpt-6-astra 两轮评审定稿,规格全文 paper-rl/RESEARCH_REVIEW_RL_REWARD.md §2):
  终局奖励一字不动;每条轨迹最多选 1 个「之后再没被改过」的编辑块 e,把它从最终补丁里剔掉
  (块触及的路径恢复到块开始前的内容),在**新开的任务环境**里再判一次,效用差
      D_e = U(P) − U(P^{-e}),   U = Y + w_f·F − w_g·G
  只加到产生这个块的那几轮生成 token 的优势上(localized, M);对照臂把同样的信号均匀摊到整条轨迹(broadcast, B),
  两者带符号总质量与绝对总质量相同,只差落点。

这个文件不碰 SkyRL 训练器;训练器侧在 credit_trainer.py,agent 侧的快照在 agents/mini_swe_tito.py。
"""

from __future__ import annotations

import hashlib
import json
import re
import shlex
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable, Optional

# 与 tito_harbor_generator._TEST_PATH_RE 同一条规则(改测试文件的块不合格;全臂一致)
TEST_PATH_RE = re.compile(r"(^|/)(tests?|testing)(/|_)|(^|/)test_[^/]*\.py$|_tests?\.py$|(^|/)conftest\.py$")


def is_test_path(path: str) -> bool:
    return bool(TEST_PATH_RE.search(path))


# ---------------------------------------------------------------------- #
# 编辑块(episode)划分与合格判定
# ---------------------------------------------------------------------- #
@dataclass
class TurnSnapshot:
    """一轮 assistant/tool 之后的仓库快照:tree = git write-tree 的树哈希,changed = 相对上一轮变化的路径(W_t)。"""

    step: int
    tree: str
    changed: list[str] = field(default_factory=list)


@dataclass
class EditEpisode:
    idx: int
    start_step: int
    end_step: int
    tree_before: str
    tree_after: str
    paths: list[str]
    eligible: bool = True
    reason: str = ""


def build_episodes(snapshots: list[TurnSnapshot]) -> list[EditEpisode]:
    """按顺序把 W_t≠∅ 的最大连续轮段切成编辑块;快照列表必须按 step 递增、且每一轮都有一条(没执行命令的轮 changed=[])。

    合格条件(评审冻结):
      1. 净效应非零(tree_before != tree_after);
      2. 不触及测试文件;
      3. 之后任何一轮都没再改过块触及的路径(改了又还原也算改过)。
    """
    if not snapshots:
        return []
    episodes: list[EditEpisode] = []
    cur: list[TurnSnapshot] = []
    prev_tree: Optional[str] = None
    i = 0
    while i <= len(snapshots):
        s = snapshots[i] if i < len(snapshots) else None
        if s is not None and s.changed:
            cur.append(s)
        else:
            if cur:
                before = prev_tree_of(snapshots, cur[0])
                paths = sorted({p for t in cur for p in t.changed})
                episodes.append(
                    EditEpisode(
                        idx=len(episodes), start_step=cur[0].step, end_step=cur[-1].step,
                        tree_before=before, tree_after=cur[-1].tree, paths=paths,
                    )
                )
                cur = []
        i += 1
    # 合格判定
    for e in episodes:
        if e.tree_before == e.tree_after:
            e.eligible, e.reason = False, "net_zero"
            continue
        if any(is_test_path(p) for p in e.paths):
            e.eligible, e.reason = False, "touches_tests"
            continue
        later = {p for s in snapshots if s.step > e.end_step for p in s.changed}
        if later & set(e.paths):
            e.eligible, e.reason = False, "path_edited_later"
            continue
        e.reason = "ok"
    return episodes


def prev_tree_of(snapshots: list[TurnSnapshot], s: TurnSnapshot) -> str:
    """块开始前的树 = 该轮之前最近一条快照的 tree;第一轮之前就是初始快照(调用方把 S_0 放在 step=-1)。"""
    prev = None
    for t in snapshots:
        if t.step >= s.step:
            break
        prev = t
    if prev is None:
        raise ValueError("快照列表里没有初始快照(step=-1 的 S_0)")
    return prev.tree


def select_episode(episodes: Iterable[EditEpisode], key: str) -> Optional[EditEpisode]:
    """合格块里按确定性哈希均匀选 1 个:key 由 generator 给(任务 id | 采样序号 | 权重版本),与 verifier 结果、完成顺序无关。"""
    best, best_h = None, None
    for e in episodes:
        if not e.eligible:
            continue
        h = int(hashlib.sha256(f"{key}|{e.idx}".encode()).hexdigest(), 16)
        if best_h is None or h < best_h:
            best, best_h = e, h
    return best


# ---------------------------------------------------------------------- #
# 容器内 git 命令(agent 侧调用;这里只拼字符串,便于单测)
# ---------------------------------------------------------------------- #
SNAP_INDEX = "/tmp/.tito_snap_index"
CF_INDEX = "/tmp/.tito_cf_index"
REPO_CD = 'cd /testbed 2>/dev/null || cd /app 2>/dev/null || cd "$(git rev-parse --show-toplevel 2>/dev/null)" || exit 97'


def cmd_snapshot_init() -> str:
    """初始快照 S_0:用 harness 自己的临时 index(不动 agent 的 .git/index / HEAD / 工作树),把整个工作树(含未跟踪)写成 tree。"""
    return (
        f"{REPO_CD}; export GIT_INDEX_FILE={SNAP_INDEX}; cp .git/index \"$GIT_INDEX_FILE\" 2>/dev/null; "
        "git add -A . >/dev/null 2>&1 || exit 98; git write-tree"
    )


def cmd_snapshot_step(prev_tree: str) -> str:
    """一轮之后的快照:输出首行 `TREE <hash>`,随后 `<status>\\t<path>`(相对上一棵树,不做 rename 检测)。"""
    return (
        f"{REPO_CD}; export GIT_INDEX_FILE={SNAP_INDEX}; git add -A . >/dev/null 2>&1 || exit 98; "
        f"T=$(git write-tree) || exit 98; echo \"TREE $T\"; git diff --name-status --no-renames {shlex.quote(prev_tree)} \"$T\""
    )


def parse_snapshot_step(output: str) -> tuple[Optional[str], list[str]]:
    tree, changed = None, []
    for line in output.splitlines():
        if line.startswith("TREE "):
            tree = line[5:].strip()
        elif "\t" in line:
            _, path = line.split("\t", 1)
            changed.append(path.strip())
    return tree, changed


def cmd_ls_tree(tree: str, paths: list[str]) -> str:
    q = " ".join(shlex.quote(p) for p in paths)
    return f"{REPO_CD}; git ls-tree {shlex.quote(tree)} -- {q}"


def parse_ls_tree(output: str) -> dict[str, tuple[str, str]]:
    """`<mode> <type> <sha>\\t<path>` → {path: (mode, sha)};目录/子模块条目跳过(只恢复 blob)。"""
    out: dict[str, tuple[str, str]] = {}
    for line in output.splitlines():
        if "\t" not in line:
            continue
        meta, path = line.split("\t", 1)
        parts = meta.split()
        if len(parts) == 3 and parts[1] == "blob":
            out[path] = (parts[0], parts[2])
    return out


def cmd_build_counterfactual(s0_tree: str, final_tree: str, paths: list[str], restore: dict[str, tuple[str, str]], max_bytes: int) -> str:
    """从最终树出发,把块触及的路径恢复到块前状态(restore 里没有的路径 = 块前不存在 ⇒ 删除),写出反事实树,
    输出 `CFTREE <hash>` + 反事实补丁(相对 S_0 的 binary diff,截到 max_bytes+1 字节以便检测超限)。"""
    parts = [f"{REPO_CD}; export GIT_INDEX_FILE={CF_INDEX}; rm -f \"$GIT_INDEX_FILE\"; git read-tree {shlex.quote(final_tree)} || exit 98"]
    if paths:
        parts.append("git update-index --force-remove -- " + " ".join(shlex.quote(p) for p in paths) + " >/dev/null 2>&1")
    for p in paths:
        if p in restore:
            mode, sha = restore[p]
            parts.append(f"git update-index --add --cacheinfo {mode},{sha},{shlex.quote(p)} || exit 98")
    parts.append("CF=$(git write-tree) || exit 98; echo \"CFTREE $CF\"")
    parts.append(f"git diff --binary --no-renames {shlex.quote(s0_tree)} \"$CF\" | head -c {int(max_bytes) + 1}")
    return "; ".join(parts)


def cmd_final_patch(s0_tree: str, final_tree: str, max_bytes: int) -> str:
    return f"{REPO_CD}; git diff --binary --no-renames {shlex.quote(s0_tree)} {shlex.quote(final_tree)} | head -c {int(max_bytes) + 1}"


def parse_counterfactual(output: str) -> tuple[Optional[str], str]:
    if not output.startswith("CFTREE "):
        return None, ""
    first, _, rest = output.partition("\n")
    return first[7:].strip(), rest


# ---------------------------------------------------------------------- #
# 效用 U = Y + w_f·F − w_g·G(SWE-bench Verified 的 report.json 口径)
# ---------------------------------------------------------------------- #
@dataclass
class Utility:
    resolved: bool
    f2p_pass: int
    f2p_total: int
    p2p_nonpass: int
    p2p_total: int

    @property
    def F(self) -> float:
        return self.f2p_pass / self.f2p_total if self.f2p_total else 0.0

    @property
    def G(self) -> float:
        return 1.0 if self.p2p_nonpass > 0 else 0.0

    def value(self, w_f: float, w_g: float) -> float:
        return float(self.resolved) + w_f * self.F - w_g * self.G

    def as_dict(self) -> dict[str, Any]:
        return {"resolved": self.resolved, "f2p_pass": self.f2p_pass, "f2p_total": self.f2p_total, "p2p_nonpass": self.p2p_nonpass, "p2p_total": self.p2p_total}


def _expected_tests(task_dir: Path, key: str) -> list[str]:
    cfg = json.loads((task_dir / "tests" / "config.json").read_text())
    v = cfg.get(key) or cfg.get(key.lower()) or []
    if isinstance(v, str):
        v = json.loads(v)
    return list(dict.fromkeys(v))  # 去重保序


def parse_utility(trial_dir: Path, task_dir: Path) -> Optional[Utility]:
    """读 verifier/report.json(SWE-bench 官方 harness 形状)。缺失/不可解析 ⇒ None(不可用,不从终局分反推)。
    F 用题目固定的 F2P 名单做分母;G:任一期望 P2P 非通过(显式失败、或没出现在 success 里 = 缺失/被 skip)。"""
    try:
        report = json.loads((trial_dir / "verifier" / "report.json").read_text())
        entry = next(iter(report.values()))
        ts = entry.get("tests_status") or {}
        f2p, p2p = ts.get("FAIL_TO_PASS") or {}, ts.get("PASS_TO_PASS") or {}
        exp_f2p, exp_p2p = _expected_tests(task_dir, "FAIL_TO_PASS"), _expected_tests(task_dir, "PASS_TO_PASS")
        if not exp_f2p:
            return None
        f_ok = set(f2p.get("success") or [])
        p_ok = set(p2p.get("success") or [])
        f2p_pass = sum(1 for t in exp_f2p if t in f_ok)
        p2p_nonpass = sum(1 for t in exp_p2p if t not in p_ok)
        resolved = bool(entry.get("resolved", False))
        return Utility(resolved=resolved, f2p_pass=f2p_pass, f2p_total=len(exp_f2p), p2p_nonpass=p2p_nonpass, p2p_total=len(exp_p2p))
    except Exception:  # noqa: BLE001
        return None


# ---------------------------------------------------------------------- #
# 信用张量(M 定位 / B 广播),评审 §B4–B5 的归一化
# ---------------------------------------------------------------------- #
@dataclass
class CreditArrays:
    localized: list[float]
    broadcast: list[float]
    s: float
    L: int
    m: int
    signed_mass: float
    abs_mass: float
    capped: bool


def build_credit_arrays(
    loss_mask: list[int], spans: list[tuple[int, int]], delta: float, lam: float, cap: float
) -> CreditArrays:
    """loss_mask:response 局部坐标的 0/1;spans:选中块各轮的生成 token 区间 [start,end)(response 局部坐标,已随截断裁剪)。

      L = 有效生成 token 数;m = 块内有效生成 token 数;
      s = 1(D=0)或 min(1, cap·m / (λ|D|·L));
      M:λ·D·s·(L/m) 落在块内 token;B:λ·D·s 落在全部有效 token。两者带符号总质量 = λ·D·s·L,绝对质量也相等。
    """
    n = len(loss_mask)
    in_span = [0] * n
    for a, b in spans:
        for t in range(max(0, a), min(n, b)):
            if loss_mask[t]:
                in_span[t] = 1
    L = sum(1 for x in loss_mask if x)
    m = sum(in_span)
    local = [0.0] * n
    bcast = [0.0] * n
    # λ≤0(零系数对照)或没有可落点/没有信号 ⇒ 全 0;不能让 λ=0 走到 s 的分母(0910 独立审计发现的 ZeroDivisionError)
    if L == 0 or m == 0 or delta == 0.0 or lam <= 0.0:
        return CreditArrays(local, bcast, 1.0, L, m, 0.0, 0.0, False)
    s = min(1.0, cap * m / (lam * abs(delta) * L))
    v_loc = lam * delta * s * L / m
    v_b = lam * delta * s
    for t in range(n):
        if loss_mask[t]:
            bcast[t] = v_b
            if in_span[t]:
                local[t] = v_loc
    return CreditArrays(local, bcast, s, L, m, lam * delta * s * L, abs(lam * delta * s) * L, s < 1.0)


# ---------------------------------------------------------------------- #
# 反事实 verifier:一个只会「应用补丁」的 harbor agent + 复用任务自带的 verifier
# ---------------------------------------------------------------------- #
try:  # harbor 只在训练/评测环境里有;纯逻辑单测不需要
    from harbor.agents.base import BaseAgent as _HarborBaseAgent
except Exception:  # noqa: BLE001
    _HarborBaseAgent = object  # type: ignore[misc,assignment]


class PatchApplyAgent(_HarborBaseAgent):  # type: ignore[misc]
    """把给定的 binary diff(相对任务初始工作树)应用到新开的环境里,然后什么都不做,交给 verifier 判。"""

    SUPPORTS_ATIF = False

    def __init__(self, logs_dir: Path, model_name: str | None = None, *, patch_text: str = "", **kwargs: Any) -> None:
        super().__init__(logs_dir=logs_dir, model_name=model_name, **kwargs)
        self._patch_text = patch_text or ""

    @staticmethod
    def name() -> str:
        return "patch-apply"

    def version(self) -> str | None:
        return "cf-1"

    async def setup(self, environment) -> None:  # noqa: D401
        return None

    async def run(self, instruction: str, environment, context) -> None:
        import tempfile

        applied, output = False, ""
        if not self._patch_text.strip():
            applied, output = True, "(empty patch: nothing to apply)"
        else:
            with tempfile.NamedTemporaryFile("w", suffix=".patch", delete=False) as f:
                f.write(self._patch_text)
                local = f.name
            await environment.upload_file(local, "/tmp/cf_apply.patch")
            res = await environment.exec(
                command=f"exec 2>&1; {REPO_CD}; git apply --binary --whitespace=nowarn /tmp/cf_apply.patch && echo CF_APPLY_OK",
                timeout_sec=120,
            )
            output = (res.stdout or "") + (res.stderr or "")
            applied = "CF_APPLY_OK" in output and res.return_code == 0
        context.metadata = {"cf_applied": applied, "cf_apply_output": output[-2000:], "cf_patch_chars": len(self._patch_text)}
        try:
            (self.logs_dir / "patch_apply.json").write_text(json.dumps(context.metadata, ensure_ascii=False))
        except Exception:  # noqa: BLE001
            pass


def make_counterfactual_trial_config(template: dict, task_path: str, trial_name: str, patch_text: str, agent_timeout_sec: int = 600) -> dict:
    """在 generator 的 harbor 模板上换掉 agent:环境/verifier/trials_dir 都沿用,保证判分口径与原轨迹一致。"""
    from copy import deepcopy

    cfg = deepcopy(template)
    cfg["task"] = {"path": task_path}
    cfg["trial_name"] = trial_name
    clean_agents = {
        "swe_skyrl_recipe.agents.clean_agents:CleanMiniSweTitoAgent",
        "swe_skyrl_recipe.agents.clean_agents:CleanPatchApplyAgent",
    }
    original_agent = cfg.get("agent", {}).get("import_path")
    cfg["agent"] = {
        "import_path": (
            "swe_skyrl_recipe.agents.clean_agents:CleanPatchApplyAgent"
            if original_agent in clean_agents else "swe_skyrl_recipe.edit_credit:PatchApplyAgent"
        ),
        "override_timeout_sec": agent_timeout_sec,
        "override_setup_timeout_sec": 600,
        "kwargs": {"patch_text": patch_text},
    }
    cfg["download_agent_logs"] = False
    return cfg