#!/usr/bin/env python3
"""Build a pi-format Qwen3.5 tool-calling SFT dataset from public traces.

Converts heterogeneous agent trajectories into OpenAI-style messages + tools
that prime-rl SFT (qwen3.5 renderer) can train on. Drops anything that
mentions a terminal-bench-2 or swe-bench-verified instance id.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import re
import sys
from pathlib import Path
from typing import Any, Iterable

from datasets import Dataset, load_dataset

WS = Path(os.environ.get("AGENTPTB_WORKSPACE", "/mnt/pvc/users/simon/agentptb/runs/grok/workspace"))
HF_HUB = Path(os.environ.get("HF_HUB_CACHE", os.path.expanduser("~/.cache/huggingface/hub")))
TASKSETS = Path("/var/lib/agentptb-cache/tb1/shared/tasksets")

PI_TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "bash",
            "description": "Execute a bash command in the current working directory. Returns stdout and stderr.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "Bash command to execute"},
                    "timeout": {"type": "number", "description": "Timeout in seconds (optional)"},
                },
                "required": ["command"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read",
            "description": "Read the contents of a file. Use offset/limit for large files.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Path to the file to read"},
                    "offset": {"type": "number", "description": "Line number to start from (1-indexed)"},
                    "limit": {"type": "number", "description": "Maximum number of lines to read"},
                },
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write",
            "description": "Write content to a file. Creates parent directories. Overwrites if it exists.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Path to the file to write"},
                    "content": {"type": "string", "description": "Content to write to the file"},
                },
                "required": ["path", "content"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "edit",
            "description": "Edit a file with exact text replacement. oldText must be unique.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Path to the file to edit"},
                    "edits": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "oldText": {"type": "string"},
                                "newText": {"type": "string"},
                            },
                            "required": ["oldText", "newText"],
                        },
                    },
                },
                "required": ["path", "edits"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "grep",
            "description": "Search file contents for a pattern.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pattern": {"type": "string"},
                    "path": {"type": "string"},
                    "glob": {"type": "string"},
                    "ignoreCase": {"type": "boolean"},
                    "literal": {"type": "boolean"},
                    "context": {"type": "number"},
                    "limit": {"type": "number"},
                },
                "required": ["pattern"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "find",
            "description": "Search for files by glob pattern.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pattern": {"type": "string"},
                    "path": {"type": "string"},
                    "limit": {"type": "number"},
                },
                "required": ["pattern"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "ls",
            "description": "List directory contents.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string"},
                    "limit": {"type": "number"},
                },
            },
        },
    },
]

SYSTEM_PROMPT = """You are an expert coding assistant operating inside pi, a coding agent harness. You help users by reading files, executing commands, editing code, and writing new files.

Available tools:
- bash: Execute a bash command in the current working directory
- read: Read the contents of a file
- write: Write content to a file
- edit: Edit a file using exact text replacement
- grep: Search file contents for a pattern
- find: Search for files by glob pattern
- ls: List directory contents

Guidelines:
- Use tools to inspect the environment before changing anything
- Prefer specialized tools (read/edit/grep/find/ls) when they fit; use bash for everything else
- Be concise. Do not narrate at length. Call a tool or give a short final answer and stop
- After a tool result, decide the next action immediately
- When the task is done, state what you did in one or two sentences and stop. Do not keep generating
- Never invent file contents; read them first
- For software bugs: reproduce, locate, patch with edit, then run the relevant tests
- For terminal tasks: inspect the workspace, run the required commands, verify outputs, then stop
"""


def load_eval_needles() -> list[str]:
    needles: set[str] = set()
    for suite in ("terminal-bench-2", "swe-bench-verified"):
        d = TASKSETS / suite
        if not d.is_dir():
            continue
        for name in os.listdir(d):
            if name.startswith("."):
                continue
            needles.add(name)
            needles.add(name.lower())
            # SWE ids also appear with slashes / dots
            needles.add(name.replace("__", "/"))
            needles.add(name.replace("-", "_"))
    # keep only reasonably specific strings
    return [n for n in sorted(needles) if len(n) >= 8]


EVAL_NEEDLES = load_eval_needles()
EVAL_NEEDLES_L = tuple(n.lower() for n in EVAL_NEEDLES)


def contaminated(text: str) -> bool:
    if not text:
        return False
    low = text.lower()
    return any(n in low for n in EVAL_NEEDLES_L)


def clip(s: str, n: int = 8000) -> str:
    if s is None:
        return ""
    s = str(s)
    if len(s) <= n:
        return s
    return s[: n // 2] + "\n...[truncated]...\n" + s[-n // 2 :]


def new_id(i: int) -> str:
    return f"call_{i}"


def assistant_with_tools(text: str, calls: list[dict[str, Any]]) -> dict[str, Any]:
    msg: dict[str, Any] = {"role": "assistant", "content": (text or "").strip()}
    if calls:
        msg["tool_calls"] = calls
    return msg


def tool_call(cid: str, name: str, args: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": cid,
        "type": "function",
        "function": {"name": name, "arguments": json.dumps(args, ensure_ascii=False)},
    }


def tool_result(cid: str, content: str) -> dict[str, Any]:
    return {"role": "tool", "tool_call_id": cid, "content": clip(content, 6000)}


def pack(messages: list[dict[str, Any]], source: str) -> dict[str, Any] | None:
    if not messages or not any(m.get("role") == "assistant" for m in messages):
        return None
    blob = json.dumps(messages, ensure_ascii=False)
    if contaminated(blob):
        return None
    # require at least one tool call or a short final answer (teaches stop)
    has_tools = any(m.get("tool_calls") for m in messages)
    last_as = next((m for m in reversed(messages) if m.get("role") == "assistant"), None)
    if not has_tools and last_as and len(str(last_as.get("content") or "")) > 2000:
        return None
    return {
        "messages": messages,
        "tools": PI_TOOLS,
        "source": source,
    }


# ---------------------------------------------------------------------------
# OpenThoughts Terminus-2
# ---------------------------------------------------------------------------

TASK_SPLIT_RE = re.compile(r"Task Description:\s*", re.IGNORECASE)
JSON_RE = re.compile(r"\{[\s\S]*\}")


def extract_task_text(user_content: str) -> str:
    parts = TASK_SPLIT_RE.split(user_content, maxsplit=1)
    body = parts[-1].strip()
    # drop trailing empty terminal screen
    body = re.split(r"\nCurrent terminal state:", body, maxsplit=1)[0].strip()
    body = re.sub(r"\n{3,}", "\n\n", body)
    return body[:8000]


def parse_terminus_json(text: str) -> dict[str, Any] | None:
    text = (text or "").strip()
    if not text:
        return None
    candidates = [text]
    fence = re.search(r"```(?:json)?\s*([\s\S]*?)```", text)
    if fence:
        candidates.insert(0, fence.group(1).strip())
    m = JSON_RE.search(text)
    if m:
        candidates.append(m.group(0))
    for c in candidates:
        try:
            obj = json.loads(c)
            if isinstance(obj, dict):
                return obj
        except Exception:
            continue
        # trailing junk: trim to last }
        try:
            end = c.rfind("}")
            if end > 0:
                obj = json.loads(c[: end + 1])
                if isinstance(obj, dict):
                    return obj
        except Exception:
            continue
    return None


def keystrokes_to_command(ks: str) -> str | None:
    if ks is None:
        return None
    ks = str(ks)
    if not ks.strip():
        return None
    # tmux-style specials we cannot map cleanly
    if re.search(r"\bC-[a-zA-Z]\b", ks) and len(ks.strip()) < 8:
        return None
    # drop trailing enter
    if ks.endswith("\n"):
        ks = ks[:-1]
    ks = ks.replace("\r\n", "\n").replace("\r", "\n")
    return ks.strip() or None


def convert_ot_row(row: dict[str, Any]) -> dict[str, Any] | None:
    convs = row.get("conversations") or []
    if len(convs) < 2:
        return None
    messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    first_user = None
    cid = 0
    pending_tool_ids: list[str] = []

    for i, turn in enumerate(convs):
        role = (turn.get("role") or turn.get("from") or "").lower()
        content = turn.get("content") or ""
        if role in ("human", "user"):
            if first_user is None:
                first_user = extract_task_text(content)
                if not first_user:
                    return None
                messages.append({"role": "user", "content": first_user})
            else:
                # terminal dump → tool results for pending calls
                dump = content
                dump = re.sub(r"^Current terminal state:\s*", "", dump)
                dump = clip(dump, 6000)
                if pending_tool_ids:
                    # attach full dump to last call; brief ack on earlier
                    for j, tid in enumerate(pending_tool_ids):
                        messages.append(tool_result(tid, dump if j == len(pending_tool_ids) - 1 else "ok"))
                    pending_tool_ids = []
                else:
                    messages.append({"role": "user", "content": dump})
        elif role in ("assistant", "gpt", "model"):
            obj = parse_terminus_json(content)
            cmds: list[str] = []
            analysis = ""
            complete = False
            if obj:
                analysis = " ".join(
                    p for p in (obj.get("analysis") or "", obj.get("plan") or "") if p
                ).strip()
                analysis = clip(analysis, 800)
                complete = bool(obj.get("task_complete"))
                for c in obj.get("commands") or []:
                    if not isinstance(c, dict):
                        continue
                    cmd = keystrokes_to_command(c.get("keystrokes", ""))
                    if cmd:
                        cmds.append(cmd)
            else:
                # non-JSON final commentary
                analysis = clip(content.strip(), 800)
                complete = True

            if cmds:
                # one bash call with the batch (pi bash is a real shell)
                joined = "\n".join(cmds)
                tid = new_id(cid)
                cid += 1
                text = analysis or "Running the next commands."
                messages.append(assistant_with_tools(text, [tool_call(tid, "bash", {"command": joined})]))
                pending_tool_ids = [tid]
            else:
                final = analysis or "Done."
                if complete or not pending_tool_ids:
                    messages.append({"role": "assistant", "content": final[:600]})
                pending_tool_ids = []

    # drop trailing unmatched tool calls (no result)
    if pending_tool_ids:
        # keep them only if we have a later user dump already consumed
        pass

    # must end on assistant text or a complete tool cycle
    if messages[-1]["role"] == "tool":
        messages.append({"role": "assistant", "content": "Task complete."})
    if messages[-1]["role"] != "assistant":
        return None
    if not any(m.get("tool_calls") for m in messages):
        return None
    return pack(messages, source=f"ot:{row.get('task', 'unknown')}")


# ---------------------------------------------------------------------------
# Real pi sessions
# ---------------------------------------------------------------------------

PI_TOOL_NAMES = {"bash", "read", "write", "edit", "grep", "find", "ls"}


def _pi_leaf_path(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Walk parentId links from the last message back to root (cycle-safe)."""
    by_id: dict[str, dict[str, Any]] = {}
    for e in entries:
        eid = e.get("id")
        if eid:
            by_id[eid] = e
    msgs = [e for e in entries if e.get("type") == "message" and e.get("id")]
    if not msgs:
        return [e for e in entries if e.get("type") == "message"]
    path: list[dict[str, Any]] = []
    seen: set[str] = set()
    cur: dict[str, Any] | None = msgs[-1]
    while cur is not None:
        eid = cur.get("id")
        if not eid or eid in seen:
            break
        seen.add(eid)
        path.append(cur)
        pid = cur.get("parentId")
        cur = by_id.get(pid) if pid else None
    path.reverse()
    return [e for e in path if e.get("type") == "message"]


def convert_pi_session(path: Path) -> dict[str, Any] | None:
    entries: list[dict[str, Any]] = []
    try:
        with path.open("r", encoding="utf-8", errors="replace") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entries.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    except OSError:
        return None
    msgs_entries = _pi_leaf_path(entries)
    if not msgs_entries:
        return None

    messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    saw_user = False
    n_tools = 0
    cid = 0
    last_call_ids: list[str] = []

    for e in msgs_entries:
        m = e.get("message") or e
        role = m.get("role")
        if role == "user":
            content = m.get("content")
            if isinstance(content, list):
                text = "\n".join(
                    p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"
                )
            else:
                text = str(content or "")
            text = text.strip()
            if not text:
                continue
            # skip slash-commands / meta
            if text.startswith("/") and len(text) < 80:
                continue
            messages.append({"role": "user", "content": clip(text, 4000)})
            saw_user = True
        elif role == "assistant":
            content = m.get("content") or []
            if isinstance(content, str):
                content = [{"type": "text", "text": content}]
            texts: list[str] = []
            calls: list[dict[str, Any]] = []
            last_call_ids = []
            for p in content:
                if not isinstance(p, dict):
                    continue
                typ = p.get("type")
                if typ in ("text",):
                    t = (p.get("text") or "").strip()
                    if t:
                        texts.append(t)
                elif typ in ("toolCall", "tool_call"):
                    name = p.get("name") or (p.get("function") or {}).get("name")
                    args = p.get("arguments") or (p.get("function") or {}).get("arguments") or {}
                    if isinstance(args, str):
                        try:
                            args = json.loads(args)
                        except json.JSONDecodeError:
                            args = {"command": args} if name == "bash" else {}
                    if name not in PI_TOOL_NAMES:
                        # map common aliases
                        alias = {
                            "shell": "bash",
                            "run": "bash",
                            "exec": "bash",
                            "cat": "read",
                            "view": "read",
                            "str_replace": "edit",
                            "str_replace_editor": "edit",
                            "create": "write",
                            "search": "grep",
                            "glob": "find",
                        }.get(name or "", "")
                        if not alias:
                            continue
                        name = alias
                        if name == "bash" and "command" not in args:
                            args = {"command": args.get("cmd") or args.get("input") or ""}
                    # normalize edit
                    if name == "edit":
                        if "edits" not in args:
                            old = args.get("oldText") or args.get("old_str") or args.get("old_string")
                            new = args.get("newText") or args.get("new_str") or args.get("new_string")
                            path_ = args.get("path") or args.get("file") or args.get("file_path")
                            if old is None or new is None or not path_:
                                continue
                            args = {"path": path_, "edits": [{"oldText": old, "newText": new}]}
                    if name == "read":
                        path_ = args.get("path") or args.get("file") or args.get("file_path")
                        if not path_:
                            continue
                        args = {k: args[k] for k in ("path", "offset", "limit") if k in args or k == "path"}
                        args["path"] = path_
                    if name == "write":
                        path_ = args.get("path") or args.get("file") or args.get("file_path")
                        content_ = args.get("content") or args.get("new_str") or args.get("file_text")
                        if not path_ or content_ is None:
                            continue
                        args = {"path": path_, "content": content_}
                    if name == "bash":
                        cmd = args.get("command") or args.get("cmd")
                        if not cmd:
                            continue
                        args = {"command": cmd}
                        if args.get("timeout"):
                            args["timeout"] = args["timeout"]
                    tid = p.get("id") or new_id(cid)
                    cid += 1
                    calls.append(tool_call(tid, name, args))
                    last_call_ids.append(tid)
                    n_tools += 1
            text = clip("\n".join(texts), 800)
            if not text and not calls:
                continue
            messages.append(assistant_with_tools(text, calls))
        elif role in ("toolResult", "tool"):
            content = m.get("content")
            if isinstance(content, list):
                text = "\n".join(
                    p.get("text", "") for p in content if isinstance(p, dict) and "text" in p
                )
            else:
                text = str(content or "")
            tid = m.get("toolCallId") or m.get("tool_call_id")
            if not tid and last_call_ids:
                tid = last_call_ids.pop(0)
            if not tid:
                messages.append({"role": "user", "content": clip(text, 4000)})
            else:
                messages.append(tool_result(tid, text))

    if not saw_user or n_tools == 0:
        return None
    if messages[-1]["role"] == "tool":
        messages.append({"role": "assistant", "content": "Done."})
    if messages[-1]["role"] != "assistant":
        return None
    return pack(messages, source=f"pi:{path.name}")


# ---------------------------------------------------------------------------
# Generic OpenHands / SWE-agent style
# ---------------------------------------------------------------------------

OH_BASH = {"execute_bash", "run_shell_command", "bash", "terminal", "run", "shell"}
OH_READ = {"read", "read_file", "open_file", "view", "cat"}
OH_WRITE = {"write", "write_file", "create_file", "create"}
OH_EDIT = {"str_replace", "str_replace_editor", "edit", "edit_file", "replace"}
OH_GREP = {"grep", "search_content", "search"}
OH_FIND = {"find", "glob", "search_files"}
OH_LS = {"ls", "list_dir", "list_files"}


def _map_oh_tool(name: str, args: dict[str, Any]) -> tuple[str, dict[str, Any]] | None:
    n = (name or "").lower()
    if n in OH_BASH:
        cmd = args.get("command") or args.get("cmd") or args.get("input")
        if not cmd:
            return None
        return "bash", {"command": str(cmd)}
    if n in OH_READ:
        path = args.get("path") or args.get("file") or args.get("file_path") or args.get("filename")
        if not path:
            return None
        out = {"path": str(path)}
        if args.get("offset") is not None:
            out["offset"] = args["offset"]
        if args.get("limit") is not None:
            out["limit"] = args["limit"]
        return "read", out
    if n in OH_WRITE:
        path = args.get("path") or args.get("file") or args.get("file_path")
        content = args.get("content") or args.get("file_text") or args.get("new_str")
        if not path or content is None:
            return None
        return "write", {"path": str(path), "content": str(content)}
    if n in OH_EDIT:
        path = args.get("path") or args.get("file") or args.get("file_path")
        old = args.get("old_str") or args.get("oldText") or args.get("old_string")
        new = args.get("new_str") or args.get("newText") or args.get("new_string")
        if not path or old is None or new is None:
            return None
        return "edit", {"path": str(path), "edits": [{"oldText": str(old), "newText": str(new)}]}
    if n in OH_GREP:
        pat = args.get("pattern") or args.get("query") or args.get("regex")
        if not pat:
            return None
        out = {"pattern": str(pat)}
        if args.get("path"):
            out["path"] = args["path"]
        if args.get("glob"):
            out["glob"] = args["glob"]
        return "grep", out
    if n in OH_FIND:
        pat = args.get("pattern") or args.get("glob") or args.get("query")
        if not pat:
            return None
        out = {"pattern": str(pat)}
        if args.get("path"):
            out["path"] = args["path"]
        return "find", out
    if n in OH_LS:
        out: dict[str, Any] = {}
        if args.get("path"):
            out["path"] = args["path"]
        return "ls", out
    return None


def convert_generic_messages(raw_messages: list[Any], source: str) -> dict[str, Any] | None:
    """Best-effort conversion of chat-style agent traces with tool_calls or XML."""
    messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    cid = 0
    n_tools = 0
    saw_user = False

    for turn in raw_messages:
        if not isinstance(turn, dict):
            continue
        role = (turn.get("role") or turn.get("from") or "").lower()
        if role in ("system",):
            continue
        content = turn.get("content")
        if role in ("user", "human"):
            if isinstance(content, list):
                text = "\n".join(
                    (p.get("text") if isinstance(p, dict) else str(p)) for p in content
                )
            else:
                text = str(content or "")
            text = text.strip()
            # strip huge system-like prefixes
            if "Task Description:" in text:
                text = extract_task_text(text)
            if text:
                messages.append({"role": "user", "content": clip(text, 4000)})
                saw_user = True
        elif role in ("assistant", "gpt", "model"):
            text = ""
            if isinstance(content, str):
                text = content
            elif isinstance(content, list):
                text = "\n".join(
                    p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"
                )
            calls_in = turn.get("tool_calls") or turn.get("function_call") or []
            if isinstance(calls_in, dict):
                calls_in = [calls_in]
            mapped: list[dict[str, Any]] = []
            for tc in calls_in:
                if not isinstance(tc, dict):
                    continue
                func = tc.get("function") or tc
                name = func.get("name") or tc.get("name")
                args = func.get("arguments") or tc.get("arguments") or {}
                if isinstance(args, str):
                    try:
                        args = json.loads(args)
                    except json.JSONDecodeError:
                        args = {"command": args}
                mapped_t = _map_oh_tool(name, args if isinstance(args, dict) else {})
                if not mapped_t:
                    continue
                n, a = mapped_t
                tid = tc.get("id") or new_id(cid)
                cid += 1
                mapped.append(tool_call(tid, n, a))
                n_tools += 1
            # XML fallback: <function=name> ...
            if not mapped and isinstance(text, str) and "<function=" in text:
                for fm in re.finditer(
                    r"<function=([a-zA-Z0-9_]+)>([\s\S]*?)</function>", text
                ):
                    name, body = fm.group(1), fm.group(2)
                    args: dict[str, Any] = {}
                    for pm in re.finditer(
                        r"<parameter=([a-zA-Z0-9_]+)>\s*([\s\S]*?)\s*</parameter>", body
                    ):
                        args[pm.group(1)] = pm.group(2)
                    mapped_t = _map_oh_tool(name, args)
                    if not mapped_t:
                        continue
                    n, a = mapped_t
                    tid = new_id(cid)
                    cid += 1
                    mapped.append(tool_call(tid, n, a))
                    n_tools += 1
                text = re.sub(r"<tool_call>[\s\S]*?</tool_call>", "", text).strip()
            text = clip(text.strip(), 800)
            if text or mapped:
                messages.append(assistant_with_tools(text, mapped))
        elif role in ("tool", "function", "observation"):
            tid = turn.get("tool_call_id") or turn.get("tool_callId") or new_id(cid)
            content_s = content if isinstance(content, str) else json.dumps(content, ensure_ascii=False)
            messages.append(tool_result(tid, str(content_s)))

    if not saw_user or n_tools == 0:
        return None
    if messages[-1]["role"] == "tool":
        messages.append({"role": "assistant", "content": "Done."})
    if messages[-1]["role"] != "assistant":
        return None
    return pack(messages, source=source)


# ---------------------------------------------------------------------------
# Synthetic pi-format episodes (written here; not eval-item-conditioned)
# ---------------------------------------------------------------------------


def synthetic_examples(rng: random.Random, n: int = 1200) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []

    def add(user: str, turns: list[dict[str, Any]], tag: str) -> None:
        msgs = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}]
        msgs.extend(turns)
        packed = pack(msgs, source=f"synth:{tag}")
        if packed:
            out.append(packed)

    # 1) ls + write + confirm
    for i in range(n // 8):
        fname = f"note_{i}.txt"
        body = rng.choice(["hello world", "ok", "status=ready", "alpha beta"])
        add(
            f"Create {fname} in the current directory containing exactly: {body}",
            [
                assistant_with_tools("Checking the directory.", [tool_call("c0", "ls", {})]),
                tool_result("c0", "README.md\nsrc/\n"),
                assistant_with_tools(
                    f"Writing {fname}.",
                    [tool_call("c1", "write", {"path": fname, "content": body + "\n"})],
                ),
                tool_result("c1", f"Wrote {fname}"),
                {"role": "assistant", "content": f"Wrote {fname} with the requested contents."},
            ],
            "write-file",
        )

    # 2) read + edit
    for i in range(n // 8):
        old = f"return {i};"
        new = f"return {i + 1};"
        add(
            f"In src/util.py the function should return {i+1} instead of {i}. Fix it.",
            [
                assistant_with_tools("Reading the file.", [tool_call("c0", "read", {"path": "src/util.py"})]),
                tool_result("c0", f"def value():\n    {old}\n"),
                assistant_with_tools(
                    "Updating the return value.",
                    [
                        tool_call(
                            "c1",
                            "edit",
                            {"path": "src/util.py", "edits": [{"oldText": old, "newText": new}]},
                        )
                    ],
                ),
                tool_result("c1", "Edited src/util.py"),
                {"role": "assistant", "content": f"Updated src/util.py so it returns {i+1}."},
            ],
            "edit-return",
        )

    # 3) grep + edit
    for i in range(n // 10):
        add(
            "Replace every TODO comment in app.py with a pass statement placeholder only where the comment is the whole line.",
            [
                assistant_with_tools(
                    "Finding TODO comments.",
                    [tool_call("c0", "grep", {"pattern": "TODO", "path": "app.py"})],
                ),
                tool_result("c0", "app.py:12:    # TODO\n"),
                assistant_with_tools("Reading context.", [tool_call("c1", "read", {"path": "app.py"})]),
                tool_result("c1", "def main():\n    # TODO\n    print('hi')\n"),
                assistant_with_tools(
                    "Replacing the placeholder.",
                    [
                        tool_call(
                            "c2",
                            "edit",
                            {
                                "path": "app.py",
                                "edits": [{"oldText": "    # TODO\n", "newText": "    pass\n"}],
                            },
                        )
                    ],
                ),
                tool_result("c2", "Edited app.py"),
                {"role": "assistant", "content": "Replaced the TODO line in app.py with pass."},
            ],
            "grep-edit",
        )

    # 4) bash tests
    for i in range(n // 8):
        add(
            "Run the unit tests and fix any failure you can see in test_math.py.",
            [
                assistant_with_tools(
                    "Running tests.",
                    [tool_call("c0", "bash", {"command": "python -m pytest -q"})],
                ),
                tool_result("c0", "F\n=== FAILURES ===\ntest_math.py::test_add failed\nassert add(2,2)==5\n"),
                assistant_with_tools("Reading the test and implementation.", [
                    tool_call("c1", "read", {"path": "test_math.py"}),
                    tool_call("c2", "find", {"pattern": "**/*.py"}),
                ]),
                tool_result("c1", "from mathlib import add\n\ndef test_add():\n    assert add(2, 2) == 4\n"),
                tool_result("c2", "mathlib.py\ntest_math.py\n"),
                assistant_with_tools("Reading implementation.", [tool_call("c3", "read", {"path": "mathlib.py"})]),
                tool_result("c3", "def add(a, b):\n    return a + b + 1\n"),
                assistant_with_tools(
                    "Removing the off-by-one.",
                    [
                        tool_call(
                            "c4",
                            "edit",
                            {
                                "path": "mathlib.py",
                                "edits": [{"oldText": "    return a + b + 1\n", "newText": "    return a + b\n"}],
                            },
                        )
                    ],
                ),
                tool_result("c4", "Edited mathlib.py"),
                assistant_with_tools("Re-running tests.", [tool_call("c5", "bash", {"command": "python -m pytest -q"})]),
                tool_result("c5", ".\n1 passed\n"),
                {"role": "assistant", "content": "Fixed an off-by-one in add() and tests pass."},
            ],
            "pytest-fix",
        )

    # 5) terminal-style: inspect, run, write output
    for i in range(n // 8):
        add(
            "Count the lines in all *.log files under /var/log-like data/logs and write the total to /output/count.txt.",
            [
                assistant_with_tools("Listing logs.", [tool_call("c0", "bash", {"command": "ls -la data/logs && wc -l data/logs/*.log"})]),
                tool_result("c0", "data/logs/a.log 10\ndata/logs/b.log 7\n17 total\n"),
                assistant_with_tools(
                    "Writing the total.",
                    [tool_call("c1", "bash", {"command": "mkdir -p /output && echo 17 > /output/count.txt && wc -c /output/count.txt"})],
                ),
                tool_result("c1", "3 /output/count.txt\n"),
                {"role": "assistant", "content": "Wrote 17 to /output/count.txt."},
            ],
            "count-logs",
        )

    # 6) stop-only after a simple question (teaches EOS without tools)
    for i in range(n // 10):
        add(
            "What tool should I use to change one unique string in a file?",
            [{"role": "assistant", "content": "Use edit with a unique oldText and the replacement newText."}],
            "stop-qa",
        )

    # 7) failed command then recover
    for i in range(n // 10):
        add(
            "Install the project's Python dependencies and run scripts/check.sh.",
            [
                assistant_with_tools("Looking for dependency files.", [tool_call("c0", "ls", {})]),
                tool_result("c0", "README.md\npyproject.toml\nscripts/\n"),
                assistant_with_tools(
                    "Installing with pip.",
                    [tool_call("c1", "bash", {"command": "pip install -e . && bash scripts/check.sh"})],
                ),
                tool_result("c1", "ERROR: No module named pip\n"),
                assistant_with_tools(
                    "pip missing; using python -m pip.",
                    [tool_call("c2", "bash", {"command": "python3 -m pip install -e . && bash scripts/check.sh"})],
                ),
                tool_result("c2", "Successfully installed proj\nok\n"),
                {"role": "assistant", "content": "Installed the package with python3 -m pip and check.sh passed."},
            ],
            "recover",
        )

    # 8) find + read a config
    for i in range(n // 12):
        add(
            "Find the pytest config and tell me the test path setting.",
            [
                assistant_with_tools("Searching for pytest config files.", [tool_call("c0", "find", {"pattern": "pytest.ini"})]),
                tool_result("c0", "pytest.ini\n"),
                assistant_with_tools("Reading pytest.ini.", [tool_call("c1", "read", {"path": "pytest.ini"})]),
                tool_result("c1", "[pytest]\ntestpaths = tests\n"),
                {"role": "assistant", "content": "pytest.ini sets testpaths = tests."},
            ],
            "find-read",
        )

    rng.shuffle(out)
    return out[:n]


# ---------------------------------------------------------------------------
# Loaders
# ---------------------------------------------------------------------------

def snapshot_dir(repo: str) -> Path | None:
    name = "datasets--" + repo.replace("/", "--")
    snaps = HF_HUB / name / "snapshots"
    if not snaps.is_dir():
        return None
    cands = [p for p in snaps.iterdir() if p.is_dir()]
    if not cands:
        return None
    # prefer snapshot that actually has data files
    def score(p: Path) -> int:
        n = 0
        for root, dirs, files in os.walk(p):
            n += sum(1 for f in files if f.endswith((".parquet", ".jsonl", ".json", ".arrow")))
            if n > 5:
                break
        return n

    cands.sort(key=score, reverse=True)
    return cands[0] if score(cands[0]) > 0 else None


def parse_xml_functions(text: str) -> tuple[str, list[tuple[str, dict[str, Any]]]]:
    """Split assistant text into preamble + XML function calls."""
    calls: list[tuple[str, dict[str, Any]]] = []
    for fm in re.finditer(r"<function=([a-zA-Z0-9_]+)>([\s\S]*?)</function>", text or ""):
        name, body = fm.group(1), fm.group(2)
        args: dict[str, Any] = {}
        for pm in re.finditer(r"<parameter=([a-zA-Z0-9_]+)>\s*([\s\S]*?)\s*</parameter>", body):
            args[pm.group(1)] = pm.group(2)
        calls.append((name, args))
    preamble = re.sub(r"<function=[a-zA-Z0-9_]+>[\s\S]*?</function>", "", text or "")
    preamble = re.sub(r"</?tool_call>", "", preamble).strip()
    return preamble, calls


def map_xml_tool(name: str, args: dict[str, Any]) -> tuple[str, dict[str, Any]] | None:
    n = (name or "").lower()
    cmd = (args.get("command") or "").lower()
    if n in ("execute_bash", "bash", "run", "shell"):
        c = args.get("command") or args.get("cmd") or args.get("input")
        if c:
            return "bash", {"command": str(c)}
    if n in ("file_editor", "str_replace_editor", "editor"):
        path = args.get("path") or args.get("file_path") or args.get("file")
        if cmd in ("view", "open", "") and path:
            out: dict[str, Any] = {"path": path}
            if args.get("view_range"):
                # e.g. [1, 80]
                vr = args["view_range"].strip().strip("[]")
                parts = [p.strip() for p in vr.split(",") if p.strip()]
                if parts:
                    try:
                        out["offset"] = int(parts[0])
                        if len(parts) > 1 and parts[1] not in ("-1", "None"):
                            out["limit"] = max(1, int(parts[1]) - int(parts[0]) + 1)
                    except ValueError:
                        pass
            return "read", out
        if cmd in ("create", "write") and path and args.get("file_text") is not None:
            return "write", {"path": path, "content": args.get("file_text") or ""}
        if cmd in ("str_replace", "replace", "edit") and path:
            old = args.get("old_str") or args.get("oldText")
            new = args.get("new_str") or args.get("newText")
            if old is None or new is None:
                return None
            return "edit", {"path": path, "edits": [{"oldText": old, "newText": new}]}
        if cmd == "insert" and path and args.get("new_str") is not None:
            # cannot insert without context; fall back to bash
            return "bash", {"command": f"python3 - <<'PY'\nfrom pathlib import Path\np=Path({path!r})\nlines=p.read_text().splitlines(True)\nidx={args.get('insert_line', 0)}\nlines.insert(int(idx), {args.get('new_str')!r} if {args.get('new_str')!r}.endswith('\\n') else {args.get('new_str')!r}+'\\n')\np.write_text(''.join(lines))\nPY"}
    if n in ("search", "grep"):
        pat = args.get("search_term") or args.get("pattern") or args.get("query")
        if not pat:
            return None
        out = {"pattern": pat}
        if args.get("path"):
            out["path"] = args["path"]
        return "grep", out
    if n == "finish":
        return None
    # last resort bash mapping
    mapped = _map_oh_tool(n, args)
    return mapped


def convert_xml_row(messages_in: list[Any], source: str) -> dict[str, Any] | None:
    messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    cid = 0
    n_tools = 0
    saw_user = False
    pending: list[str] = []
    first_user = True
    for turn in messages_in:
        if not isinstance(turn, dict):
            continue
        role = (turn.get("role") or "").lower()
        content = turn.get("content") or ""
        if role == "system":
            continue
        if role == "user":
            text = content if isinstance(content, str) else str(content)
            if first_user:
                # drop the huge function-spec system dump; keep the issue
                if "<github_issue>" in text:
                    m = re.search(r"<github_issue>([\s\S]*?)</github_issue>", text)
                    text = ("Fix the following issue.\n" + m.group(1).strip()) if m else text
                elif "Task Description:" in text:
                    text = extract_task_text(text)
                else:
                    # strip leading tool spec if present
                    if "BEGIN FUNCTION" in text or "We have access to the following functions" in text:
                        parts = re.split(r"\n(?=Consider the following|Please|Your task|The issue)", text, maxsplit=1)
                        if len(parts) == 2:
                            text = parts[1]
                messages.append({"role": "user", "content": clip(text.strip(), 4000)})
                saw_user = True
                first_user = False
            else:
                dump = clip(text, 6000)
                if pending:
                    for j, tid in enumerate(pending):
                        messages.append(tool_result(tid, dump if j == len(pending) - 1 else "ok"))
                    pending = []
                else:
                    messages.append({"role": "user", "content": dump})
        elif role == "assistant":
            preamble, raw_calls = parse_xml_functions(content if isinstance(content, str) else "")
            mapped: list[dict[str, Any]] = []
            pending = []
            finish_text = ""
            for name, args in raw_calls:
                if name.lower() == "finish":
                    finish_text = args.get("thought") or args.get("summary") or preamble or "Done."
                    continue
                mt = map_xml_tool(name, args)
                if not mt:
                    continue
                n, a = mt
                tid = new_id(cid)
                cid += 1
                mapped.append(tool_call(tid, n, a))
                pending.append(tid)
                n_tools += 1
            text = clip((preamble or finish_text).strip(), 800)
            if mapped:
                messages.append(assistant_with_tools(text or "Working on it.", mapped))
            elif finish_text or preamble:
                messages.append({"role": "assistant", "content": clip(finish_text or preamble, 600)})
                pending = []
    if not saw_user or n_tools == 0:
        return None
    if messages[-1]["role"] == "tool":
        messages.append({"role": "assistant", "content": "Done."})
    if messages[-1]["role"] != "assistant":
        return None
    return pack(messages, source=source)


def load_ot_like(repo: str, tag: str, max_rows: int | None = None, skip_math: bool = False) -> Iterable[dict[str, Any]]:
    p = snapshot_dir(repo)
    if p is None:
        print(f"{tag} not on disk", file=sys.stderr)
        return
    files = sorted(p.rglob("*.parquet"))
    if skip_math:
        files = [f for f in files if "math" not in f.as_posix()]
    n_ok = n_bad = n_seen = 0
    for f in files:
        try:
            ds = load_dataset("parquet", data_files=str(f), split="train")
        except Exception as e:
            print(f"{tag} skip file {f.name}: {e}", file=sys.stderr)
            continue
        for row in ds:
            if max_rows is not None and n_ok >= max_rows:
                print(f"{tag} converted {n_ok} skipped {n_bad} (capped)", flush=True)
                return
            n_seen += 1
            if n_seen % 2000 == 0:
                print(f"{tag} progress {n_seen} ok {n_ok} bad {n_bad}", flush=True)
            rec = convert_ot_row(row)
            if rec:
                rec["source"] = f"{tag}:{row.get('task', n_seen)}"
                n_ok += 1
                yield rec
            else:
                n_bad += 1
    print(f"{tag} converted {n_ok} skipped {n_bad} seen {n_seen}", flush=True)


def load_ot() -> Iterable[dict[str, Any]]:
    yield from load_ot_like("open-thoughts/OpenThoughts-Agent-v1-SFT", "ot")


def load_pi_traces() -> Iterable[dict[str, Any]]:
    p = snapshot_dir("MaxDevv/real-pi-coding-agent-traces-sessions")
    if p is None:
        print("pi traces not on disk yet", file=sys.stderr)
        return
    files = [f for f in p.rglob("*.jsonl") if f.name != "manifest.jsonl"]
    print(f"pi session files {len(files)}", flush=True)
    n_ok = n_bad = 0
    for f in files:
        try:
            rec = convert_pi_session(f)
        except Exception:
            n_bad += 1
            continue
        if rec:
            n_ok += 1
            yield rec
        else:
            n_bad += 1
    print(f"pi converted {n_ok} skipped {n_bad}", flush=True)


def _iter_hf_rows(repo: str, max_rows: int | None = None):
    p = snapshot_dir(repo)
    if p is None:
        print(f"{repo} not on disk", file=sys.stderr)
        return
    parq = list(p.rglob("*.parquet"))
    jsonl = list(p.rglob("*.jsonl"))
    rows = 0
    if parq:
        try:
            ds = load_dataset("parquet", data_files=[str(f) for f in parq], split="train")
        except Exception as e:
            print(f"failed parquet {repo}: {e}", file=sys.stderr)
            ds = None
        if ds is not None:
            for row in ds:
                yield row
                rows += 1
                if max_rows and rows >= max_rows:
                    return
            return
    for f in jsonl:
        with f.open("r", encoding="utf-8", errors="replace") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    yield json.loads(line)
                except json.JSONDecodeError:
                    continue
                rows += 1
                if max_rows and rows >= max_rows:
                    return


def _row_messages(row: dict[str, Any]) -> list[Any] | None:
    for key in ("messages", "conversations", "conversation", "trajectory", "turns"):
        v = row.get(key)
        if isinstance(v, list) and v:
            return v
        if isinstance(v, str):
            try:
                obj = json.loads(v)
                if isinstance(obj, list) and obj:
                    return obj
            except json.JSONDecodeError:
                pass
    return None


def load_generic(repo: str, tag: str, max_rows: int | None = None) -> Iterable[dict[str, Any]]:
    n_ok = n_bad = 0
    for i, row in enumerate(_iter_hf_rows(repo, max_rows=max_rows)):
        msgs = _row_messages(row)
        if not msgs:
            n_bad += 1
            continue
        rec = convert_generic_messages(msgs, source=f"{tag}:{i}")
        if rec:
            n_ok += 1
            yield rec
        else:
            n_bad += 1
    print(f"{tag} converted {n_ok} skipped {n_bad}", flush=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--synth", type=int, default=1500)
    ap.add_argument("--max-generic", type=int, default=40000)
    args = ap.parse_args()

    rng = random.Random(args.seed)
    records: list[dict[str, Any]] = []

    print("building synthetic...", flush=True)
    records.extend(synthetic_examples(rng, args.synth))
    print(f"synth {len(records)}", flush=True)

    for rec in load_ot():
        records.append(rec)
    for rec in load_pi_traces():
        records.append(rec)

    # Terminus-2 style public traces (not eval items)
    for rec in load_ot_like(
        "open-thoughts/OpenThoughts-Agent-SFT-100K", "ot100k", max_rows=min(args.max_generic, 25000)
    ):
        records.append(rec)
    for rec in load_ot_like(
        "nvidia/Nemotron-Terminal-Corpus",
        "nemo-term",
        max_rows=min(args.max_generic, 20000),
        skip_math=True,
    ):
        records.append(rec)

    # XML function-call SWE traces
    r2e = snapshot_dir("R2E-Gym/R2EGym-SFT-Trajectories")
    if r2e:
        n_ok = n_bad = 0
        files = [str(f) for f in r2e.rglob("*.parquet")]
        if files:
            ds = load_dataset("parquet", data_files=files, split="train")
            for i, row in enumerate(ds):
                rec = convert_xml_row(row.get("messages") or [], source=f"r2e:{i}")
                if rec:
                    n_ok += 1
                    records.append(rec)
                else:
                    n_bad += 1
        print(f"r2e converted {n_ok} skipped {n_bad}", flush=True)
    sweg = snapshot_dir("SWE-Gym/OpenHands-SFT-Trajectories")
    if sweg:
        n_ok = n_bad = 0
        files = [str(f) for f in sweg.rglob("*.parquet")]
        if files:
            ds = load_dataset("parquet", data_files=files, split="train")
            for i, row in enumerate(ds):
                rec = convert_xml_row(row.get("messages") or [], source=f"swegym:{i}")
                if rec:
                    n_ok += 1
                    records.append(rec)
                else:
                    n_bad += 1
        print(f"swegym converted {n_ok} skipped {n_bad}", flush=True)

    # de-dup by hash of assistant tool calls + last user
    seen: set[str] = set()
    uniq: list[dict[str, Any]] = []
    for rec in records:
        key = hashlib.md5(json.dumps(rec["messages"], ensure_ascii=False, sort_keys=True).encode()).hexdigest()
        if key in seen:
            continue
        seen.add(key)
        uniq.append(rec)
    rng.shuffle(uniq)
    print(f"unique records {len(uniq)}", flush=True)

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    table = {
        "messages": [r["messages"] for r in uniq],
        "tools": [json.dumps(r["tools"]) for r in uniq],
        "source": [r["source"] for r in uniq],
    }
    ds = Dataset.from_dict(table)
    pq_path = out / "train.parquet"
    ds.to_parquet(str(pq_path))
    sample_path = out / "sample.jsonl"
    with sample_path.open("w") as f:
        for r in uniq[:20]:
            f.write(json.dumps({"source": r["source"], "n_msg": len(r["messages"])}) + "\n")
    # source histogram
    from collections import Counter

    hist = Counter(r["source"].split(":")[0] for r in uniq)
    print("sources", dict(hist), flush=True)
    print(f"wrote {len(uniq)} to {pq_path}", flush=True)


if __name__ == "__main__":
    main()