#!/usr/bin/env python3
"""SFT rows: chain commands, don't repeat ls, stop after the artifact exists."""
from __future__ import annotations

import json
import random
from pathlib import Path

from datasets import Dataset, load_dataset

SYS = """You are an expert coding assistant operating inside pi.
Be concise. Prefer one bash call that chains steps with && or newlines.
Never repeat the same ls/find/pwd. After the required output exists, say one short sentence and STOP. Do not celebrate via more commands."""

TOOLS = json.dumps(
    [
        {
            "type": "function",
            "function": {
                "name": "bash",
                "description": "Execute a bash command",
                "parameters": {
                    "type": "object",
                    "properties": {"command": {"type": "string"}},
                    "required": ["command"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "write",
                "description": "Write a file",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
                    "required": ["path", "content"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "read",
                "description": "Read a file",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}},
                    "required": ["path"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "edit",
                "description": "Edit a file",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}, "edits": {"type": "array"}},
                    "required": ["path", "edits"],
                },
            },
        },
    ]
)


def tc(i, name, args):
    return {"id": f"c{i}", "type": "function", "function": {"name": name, "arguments": json.dumps(args)}}


def asst(text, calls=None):
    m = {"role": "assistant", "content": text}
    if calls:
        m["tool_calls"] = calls
    return m


def tool(i, content):
    return {"role": "tool", "tool_call_id": f"c{i}", "content": content}


def pack(msgs, src):
    return {"messages": [{"role": "system", "content": SYS}] + msgs, "tools": TOOLS, "source": src}


def main():
    rng = random.Random(11)
    rows = []

    for i in range(400):
        rows.append(
            pack(
                [
                    {"role": "user", "content": f"Write READY to /output/status{i}.txt then stop."},
                    asst(
                        "Writing the file.",
                        [tc(0, "bash", {"command": f"mkdir -p /output && printf 'READY\\n' > /output/status{i}.txt && cat /output/status{i}.txt"})],
                    ),
                    tool(0, "READY\n"),
                    asst("Wrote /output/status.txt."),
                ],
                "stop:write",
            )
        )

    for i in range(250):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "Explore once, then compile src/main.c to /output/app and stop."},
                    asst(
                        "Listing and compiling in one command.",
                        [tc(0, "bash", {"command": "ls -la && gcc -O2 src/main.c -o /output/app && ls -la /output/app"})],
                    ),
                    tool(0, "-rwxr-xr-x 1 root root 16384 app\n"),
                    asst("Built /output/app."),
                ],
                "stop:build",
            )
        )

    for i in range(250):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "Fix the off-by-one in add.py and run the test. Stop when it passes."},
                    asst("Reading, patching, testing.", [tc(0, "bash", {"command": "cat add.py && python3 - <<'PY'\nfrom pathlib import Path\np=Path('add.py')\np.write_text(p.read_text().replace('a + b + 1','a + b'))\nPY\npython3 -c 'from add import add; assert add(2,2)==4'"})]),
                    tool(0, ""),
                    asst("Fixed add.py; test passed."),
                ],
                "stop:fix",
            )
        )

    # saw listing already -> do not list again
    for i in range(200):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "There is a bug in util.py. Find it and patch it."},
                    asst("Listing once.", [tc(0, "bash", {"command": "ls && cat util.py"})]),
                    tool(0, "README.md\nutil.py\ndef inc(x):\n    return x\n"),
                    asst(
                        "Patching without listing again.",
                        [
                            tc(
                                1,
                                "edit",
                                {"path": "util.py", "edits": [{"oldText": "    return x\n", "newText": "    return x + 1\n"}]},
                            )
                        ],
                    ),
                    tool(1, "Edited util.py"),
                    asst("inc now returns x+1."),
                ],
                "stop:no-relist",
            )
        )

    # mix prior finish data
    prev = Path("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-finish/train.parquet")
    if prev.exists():
        ds = load_dataset("parquet", data_files=str(prev), split="train")
        idx = list(range(len(ds)))
        rng.shuffle(idx)
        for j in idx[:4000]:
            rec = ds[j]
            rows.append({"messages": rec["messages"], "tools": rec["tools"], "source": rec["source"]})

    rng.shuffle(rows)
    out = Path("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-stop")
    out.mkdir(parents=True, exist_ok=True)
    Dataset.from_dict(
        {
            "messages": [r["messages"] for r in rows],
            "tools": [r["tools"] if isinstance(r["tools"], str) else json.dumps(r["tools"]) for r in rows],
            "source": [r["source"] for r in rows],
        }
    ).to_parquet(str(out / "train.parquet"))
    print("wrote", len(rows), out)


if __name__ == "__main__":
    main()