#!/usr/bin/env python3
"""Prepend a real orientation turn to a fraction of trajectories.

The eval harness starts the agent in an empty scratch dir (`/workspace`) with the project
somewhere else; the raw corpora all start already inside the repo, so a model trained on them
just writes files wherever it landed. This inserts, before the first assistant turn:

    assistant: <short think> + bash "pwd; ls -a; echo ---; ls -d /testbed /app /workspace ..."
    tool:      <the real output, captured from actual task containers>

The observation is not invented — `assets/orient/*.txt` are verbatim captures from three real
SWE task images through the same broker. Run this BEFORE diversify_paths.py so the path
rewrite carries through the observation too, and the model sees roots other than /testbed.
"""

import argparse
import glob
import json
import os
import random

THINKS = [
    "First, work out where I am and where the project actually is — the shell does not start inside the repo.",
    "Before touching anything I should check the working directory and locate the project tree.",
    "Let me orient myself: print the cwd and look for the project root.",
    "I don't know which directory this shell starts in. Check that first, then find the codebase.",
    "Start by locating the repository; the current directory is probably not it.",
]

PROBE = "pwd; ls -a; echo ---; ls -d /testbed /app /workspace /repo /code /srv 2>/dev/null"

FOLLOWUPS = [
    "The project is at /testbed. I'll use absolute paths under it from here on.",
    "Right — the repo is /testbed, and the cwd is an empty scratch dir. Work in /testbed.",
    "/testbed holds the checkout. Every path from here on is absolute under /testbed.",
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--obs-dir", default="assets/orient")
    ap.add_argument("--p", type=float, default=0.40)
    ap.add_argument("--seed", type=int, default=23)
    args = ap.parse_args()
    rng = random.Random(args.seed)

    observations = []
    for f in sorted(glob.glob(args.obs_dir + "/*.txt")):
        text = open(f).read().strip()
        if text:
            observations.append(text)
    if not observations:
        raise SystemExit(f"no captured observations in {args.obs_dir}")
    print(f"{len(observations)} captured observations")

    import pyarrow.parquet as pq
    from datasets import Dataset

    rows, n_aug = [], 0
    for f in sorted(glob.glob(args.src + "/*.parquet")):
        for batch in pq.ParquetFile(f).iter_batches(batch_size=256):
            for row in batch.to_pylist():
                msgs = list(row["messages"])
                first_assist = next((i for i, m in enumerate(msgs) if m["role"] == "assistant"), None)
                uses_testbed = "/testbed" in json.dumps(msgs)
                if first_assist is not None and uses_testbed and rng.random() < args.p:
                    call_id = "chatcmpl-tool-%016x" % rng.getrandbits(56)
                    probe = [
                        {
                            "role": "assistant",
                            "reasoning_content": rng.choice(THINKS),
                            "content": "",
                            "tool_calls": [
                                {
                                    "id": call_id,
                                    "type": "function",
                                    "function": {"name": "bash", "arguments": json.dumps({"command": PROBE})},
                                }
                            ],
                        },
                        {"role": "tool", "tool_call_id": call_id, "name": "bash",
                         "content": rng.choice(observations)},
                    ]
                    # fold the conclusion into the reasoning of the original first turn
                    head = dict(msgs[first_assist])
                    head["reasoning_content"] = (
                        rng.choice(FOLLOWUPS) + " " + (head.get("reasoning_content") or "")
                    ).strip()
                    msgs = msgs[:first_assist] + probe + [head] + msgs[first_assist + 1:]
                    n_aug += 1
                    row = {**row, "messages": msgs, "n_assistant": row["n_assistant"] + 1}
                rows.append(row)

    rng.shuffle(rows)
    os.makedirs(args.out, exist_ok=True)
    ds = Dataset.from_list(rows)
    shards = max(1, len(rows) // 20000 + 1)
    for i in range(shards):
        ds.shard(num_shards=shards, index=i).to_parquet(
            os.path.join(args.out, f"train-{i:05d}-of-{shards:05d}.parquet")
        )
    print(f"wrote {len(rows)} rows ({n_aug} with an orientation turn) -> {args.out}")


if __name__ == "__main__":
    main()