code/local_train/chain_pack.py

"""SFT rows that teach the live sanity chain: edit microtarget → submit → rejected → edit → submit."""

from __future__ import annotations

import json
import random
import re
from pathlib import Path
from typing import Any

from albedo_eval_service.remote.dataset import EvalSample, apply_submit_protocol, format_messages
from albedo_eval_service.shared.observation_format import first_bash_block
from sanity_service.chain import followup_instruction, micro_instruction

from .constants import DEFAULT_DATA_ROOT, DEFAULT_PACK_DIR, KEEP_ORIGINAL_RATIO, TOKENIZER_DIR
from .pack import PackedExample, _iter_raw, gold_paths, is_edit_command, _summary
from .think import wrap_completion

_FILE_RE = re.compile(
    r"(?:[\w.-]+/){1,8}[\w.-]+\.(?:py|rs|ts|js|tsx|jsx|go|java|c|h|hpp|cpp|cc|rb|toml|cmake)"
    r"|(?:[\w.-]+/)*CMakeLists\.txt"
)
_FUNC_RES = (
    re.compile(r"Function:\s*([A-Za-z_][A-Za-z0-9_]{1,64})"),
    re.compile(r"\*\*([A-Za-z_][A-Za-z0-9_]{1,64})\*\*"),
    re.compile(r"\b(?:export\s+)?(?:async\s+)?function\s+([A-Za-z_][A-Za-z0-9_]{1,64})\s*[<(]"),
    re.compile(r"\b(?:def|fn|func)\s+([A-Za-z_][A-Za-z0-9_]{1,64})\s*[(<\[]"),
    re.compile(r"\bclass\s+([A-Za-z_][A-Za-z0-9_]{1,64})\b"),
    re.compile(r"#\s*define\s+([A-Za-z_][A-Za-z0-9_]{2,64})"),
)
_WEAK_FUNCS = {
    "",
    "the helper",
    "to",
    "the target file",
    "helper",
    "function",
    "func",
    "async",
    "await",
    "class",
    "export",
    "import",
    "const",
    "return",
    "true",
    "false",
    "undefined",
    "null",
    "this",
    "self",
    "type",
    "name",
    "value",
    "values",
    "data",
    "item",
    "index",
    "error",
    "errors",
    "message",
    "options",
    "cause",
    "payload",
    "according",
    "containing",
}
_SKIP_FILE_PARTS = ("node_modules", "newfile.py", "filename.py")
_OBS = (
    "<returncode>0</returncode>\n<output>\ncommand completed with no captured output\n</output>"
)
_REJECTION = (
    "Not quite there yet: please double-check the change against the surrounding call "
    "sites and handle any case you may have missed, then submit again the same way."
)


def infer_micro(text: str, gold: str = "") -> dict[str, str]:
    """Name a real file+symbol the way live GLM does, not English leftovers."""
    hay = f"{gold}\n{text}"
    files = _candidate_files(hay)
    funcs = _candidate_funcs(hay)
    path, func = _pair_micro(files, funcs, hay)
    name = path.rsplit("/", 1)[-1] if path else "the target file"
    request = (
        f"Make a small concrete edit in {name} ({func}) so the adjacent call sites stay consistent."
    )
    return {"file": path or name, "function": func, "request": request, "message": ""}


def _candidate_files(hay: str) -> list[str]:
    seen: list[str] = []
    for match in _FILE_RE.finditer(hay):
        path = match.group(0)
        if any(part in path for part in _SKIP_FILE_PARTS):
            continue
        if path.startswith("n/") or "/testbed/" in path:
            continue
        if path not in seen:
            seen.append(path)
    return seen


def _candidate_funcs(hay: str) -> list[str]:
    seen: list[str] = []
    for pattern in _FUNC_RES:
        for match in pattern.finditer(hay):
            name = match.group(1)
            if name.lower() in _WEAK_FUNCS or name in seen:
                continue
            seen.append(name)
    return seen


def _file_score(path: str, hay: str) -> int:
    head = hay[:8000]
    base = path.rsplit("/", 1)[-1]
    score = head.lower().count(path.lower()) * 3 + head.lower().count(base.lower())
    if "/" in path:
        score += 1
    if path.endswith("CMakeLists.txt"):
        score += head.lower().count("cmake") * 4
    return score


def _pair_micro(files: list[str], funcs: list[str], hay: str) -> tuple[str, str]:
    files = sorted(files, key=lambda path: (-_file_score(path, hay), path))
    for path in files:
        stem = path.rsplit("/", 1)[-1].rsplit(".", 1)[0]
        for func in funcs:
            if func.lower() == stem.lower() or stem.lower() == func.lower():
                return path, func
    path = files[0] if files else ""
    if funcs:
        return path, funcs[0]
    stem = path.rsplit("/", 1)[-1].rsplit(".", 1)[0] if path else ""
    if stem and stem.lower() not in _WEAK_FUNCS:
        return path, stem
    return path, "the helper"


def followup_edit_command(path: str, func: str) -> str:
    target = path or "src/main.py"
    needle = func if func and func != "the helper" else "return "
    return f"sed -i '/{needle}/a\\    # handle adjacent call sites' {target}"


def pack_chain(
    *,
    dataset_root: Path = DEFAULT_DATA_ROOT,
    out_dir: Path = DEFAULT_PACK_DIR,
    max_groups: int = 400,
    seed: str = "chain-pack",
    n_salts: int = 3,
    mix_pack: Path | None = None,
    mix_n: int = 1500,
    mix_edit_n: int = 0,
    mix_submit_n: int = 0,
    tokenizer_path: Path | None = None,
) -> Path:
    dataset_root = Path(dataset_root)
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    tokenizer = str(tokenizer_path or TOKENIZER_DIR)
    salts = [f"{seed}-{i}" for i in range(max(1, n_salts))]
    raw = [
        item
        for item in _iter_raw(dataset_root, max_raw=max(max_groups * 12, 512), seed=seed)
        if item.kind == "edit" and gold_paths(item.gold)
    ]
    rng = random.Random(seed)
    rng.shuffle(raw)
    raw = raw[: max(1, max_groups)]

    examples: list[PackedExample] = []
    for index, item in enumerate(raw):
        salt = salts[index % len(salts)]
        sample = EvalSample(
            sample_id=f"{item.shard}:{item.row}:{item.turn_idx}",
            prompt="",
            messages=list(item.messages),
        )
        rewritten = apply_submit_protocol(
            [sample],
            salt=salt,
            keep_original_ratio=KEEP_ORIGINAL_RATIO,
            tokenizer_path=tokenizer,
            enable_thinking=True,
        )[0]
        clause = rewritten.submit_command
        marker = rewritten.submit_marker
        if not clause or not marker:
            continue
        micro = infer_micro(
            "\n".join(m.get("content") or "" for m in (rewritten.messages or [])),
            item.gold,
        )
        if not micro.get("file"):
            continue
        instruction = micro_instruction(micro, clause)
        edit_bash = first_bash_block(item.gold)
        if not edit_bash or not is_edit_command(edit_bash):
            continue
        edit_completion = wrap_completion(item.gold, edit_bash)
        submit_completion = wrap_completion(
            f"THOUGHT: Work is saved. Submit with the exact command.\n```bash\n{clause}\n```",
            clause,
        )
        reject_bash = followup_edit_command(micro["file"], micro["function"])
        reject_completion = wrap_completion(
            f"THOUGHT: Reviewer asked for an adjacent-site fix before resubmitting.\n```bash\n{reject_bash}\n```",
            reject_bash,
        )
        if not (edit_completion and submit_completion and reject_completion):
            continue
        prefix = list(rewritten.messages or item.messages)
        rows = _chain_rows(
            sample_id=rewritten.sample_id,
            prefix=prefix,
            instruction=instruction,
            edit_completion=edit_completion,
            submit_completion=submit_completion,
            reject_completion=reject_completion,
            clause=clause,
            marker=marker,
            rewrite_mode=rewritten.rewrite_mode,
            item=item,
            salt=salt,
            tokenizer=tokenizer,
            gold_paths=gold_paths(item.gold),
        )
        examples.extend(rows)

    if mix_pack and Path(mix_pack).is_file():
        if mix_edit_n or mix_submit_n:
            examples.extend(_mix_stage_a(Path(mix_pack), rng, mix_edit_n, mix_submit_n))
        elif mix_n:
            examples.extend(_mix_balanced(Path(mix_pack), mix_n, rng))

    pack_path = out_dir / f"sft-chain-{seed}.jsonl"
    with pack_path.open("w") as handle:
        for example in examples:
            handle.write(json.dumps(example.as_dict(), ensure_ascii=False) + "\n")
    summary = _summary(examples)
    summary["chain_groups"] = sum(1 for e in examples if e.kind.startswith("chain_"))
    (out_dir / f"sft-chain-{seed}.meta.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2), flush=True)
    print(f"pack: {pack_path}", flush=True)
    return pack_path


def _chain_rows(
    *,
    sample_id: str,
    prefix: list[dict[str, str]],
    instruction: str,
    edit_completion: str,
    submit_completion: str,
    reject_completion: str,
    clause: str,
    marker: str,
    rewrite_mode: str,
    item: Any,
    salt: str,
    tokenizer: str,
    gold_paths: list[str],
) -> list[PackedExample]:
    micro_msgs = prefix + [{"role": "user", "content": instruction}]
    after_edit = micro_msgs + [
        {"role": "assistant", "content": edit_completion},
        {"role": "user", "content": _OBS},
    ]
    after_submit = after_edit + [
        {"role": "assistant", "content": submit_completion},
        {"role": "user", "content": followup_instruction(_REJECTION, clause, first=False)},
    ]
    after_reject_edit = after_submit + [
        {"role": "assistant", "content": reject_completion},
        {"role": "user", "content": _OBS},
    ]
    after_second_submit = after_reject_edit + [
        {"role": "assistant", "content": submit_completion},
        {
            "role": "user",
            "content": followup_instruction(
                "Thanks, received. Continue with the original issue; do not submit again until you make a new edit.",
                clause,
                first=True,
            ),
        },
    ]
    specs = [
        (f"{sample_id}:chain_micro_edit", micro_msgs, edit_completion, "chain_micro_edit", "at_edit"),
        (f"{sample_id}:chain_micro_submit", after_edit, submit_completion, "chain_micro_submit", "post_edit"),
        (f"{sample_id}:chain_reject_edit", after_submit, reject_completion, "chain_reject_edit", "at_edit"),
        (f"{sample_id}:chain_reject_submit", after_reject_edit, submit_completion, "chain_reject_submit", "post_edit"),
        (f"{sample_id}:chain_followup_edit", after_second_submit, edit_completion, "chain_followup_edit", "at_edit"),
    ]
    out: list[PackedExample] = []
    for sid, messages, completion, kind, phase in specs:
        out.append(
            PackedExample(
                sample_id=sid,
                prompt=format_messages(messages, tokenizer_path=tokenizer, enable_thinking=True),
                completion=completion,
                source=item.source,
                phase=phase,
                kind=kind,
                family=item.family,
                language=item.language,
                repo=item.repo,
                submit_command=clause,
                submit_marker=marker,
                rewrite_mode=rewrite_mode,
                gold_paths=gold_paths,
                salt=salt,
            )
        )
    return out


def _row_to_example(row: dict[str, Any]) -> PackedExample:
    return PackedExample(
        sample_id=str(row.get("sample_id") or ""),
        prompt=str(row.get("prompt") or ""),
        completion=str(row.get("completion") or ""),
        source=str(row.get("source") or ""),
        phase=str(row.get("phase") or "explore"),
        kind=str(row.get("kind") or "explore"),
        family=str(row.get("family") or ""),
        language=str(row.get("language") or ""),
        repo=str(row.get("repo") or ""),
        submit_command=str(row.get("submit_command") or ""),
        submit_marker=str(row.get("submit_marker") or ""),
        rewrite_mode=str(row.get("rewrite_mode") or ""),
        gold_paths=list(row.get("gold_paths") or []),
        salt=str(row.get("salt") or ""),
    )


def _take(rows: list[dict[str, Any]], n: int, rng: random.Random) -> list[PackedExample]:
    rng.shuffle(rows)
    return [_row_to_example(row) for row in rows[: max(0, n)]]


def _mix_balanced(path: Path, n: int, rng: random.Random) -> list[PackedExample]:
    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    keep = [row for row in rows if row.get("kind") in {"edit", "explore"}]
    return _take(keep, n, rng)


def _mix_stage_a(
    path: Path, rng: random.Random, edit_n: int, submit_n: int
) -> list[PackedExample]:
    """Official gold mix: at_edit first, then any edit; post_edit submit, then any submit."""
    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    at_edit = [row for row in rows if row.get("kind") == "edit" and row.get("phase") == "at_edit"]
    other_edit = [
        row
        for row in rows
        if row.get("kind") == "edit" and row.get("phase") != "at_edit"
    ]
    post_submit = [
        row for row in rows if row.get("kind") == "submit" and row.get("phase") == "post_edit"
    ]
    other_submit = [
        row
        for row in rows
        if row.get("kind") == "submit" and row.get("phase") != "post_edit"
    ]
    edits = _take(at_edit, edit_n, rng)
    if len(edits) < edit_n:
        edits.extend(_take(other_edit, edit_n - len(edits), rng))
    submits = _take(post_submit, submit_n, rng)
    if len(submits) < submit_n:
        submits.extend(_take(other_submit, submit_n - len(submits), rng))
    return edits + submits