#!/usr/bin/env python3
"""Synthetic finish-the-task + no-loop trajectories in pi tool format."""
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, a coding agent harness. You help users by reading files, executing commands, editing code, and writing new files.

Guidelines:
- Be concise. Call a tool or give a short final answer and stop.
- Chain shell steps in one bash call. Do not loop ls/cd.
- When the required output exists, stop.
"""

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": "ls",
                "description": "List a directory",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}},
                },
            },
        },
        {
            "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(2)
    rows = []

    for i in range(400):
        fname = f"out_{i}.txt"
        body = rng.choice(["ok", "done", "42", "hello", "pass"])
        rows.append(
            pack(
                [
                    {"role": "user", "content": f"Write {body} to /output/{fname} and stop."},
                    asst("Writing the file.", [tc(0, "bash", {"command": f"mkdir -p /output && printf '%s\\n' '{body}' > /output/{fname} && wc -c /output/{fname}"})]),
                    tool(0, f"4 /output/{fname}\n"),
                    asst(f"Wrote /output/{fname}."),
                ],
                "fin:write",
            )
        )

    for i in range(300):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "Explore the repo and set the version in pkg/__init__.py to 1.2.3."},
                    asst("Listing once, then editing.", [tc(0, "bash", {"command": "ls && sed -n '1,20p' pkg/__init__.py"})]),
                    tool(0, "__version__ = '0.1.0'\n"),
                    asst(
                        "Updating the version.",
                        [
                            tc(
                                1,
                                "edit",
                                {
                                    "path": "pkg/__init__.py",
                                    "edits": [{"oldText": "__version__ = '0.1.0'", "newText": "__version__ = '1.2.3'"}],
                                },
                            )
                        ],
                    ),
                    tool(1, "Edited pkg/__init__.py"),
                    asst("Set __version__ to 1.2.3."),
                ],
                "fin:edit-once",
            )
        )

    for i in range(250):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "Run the tests. If they pass, stop. Do not keep listing files."},
                    asst("Running tests.", [tc(0, "bash", {"command": "python -m pytest -q"})]),
                    tool(0, ".\n1 passed\n"),
                    asst("Tests passed."),
                ],
                "fin:tests-pass",
            )
        )

    for i in range(250):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "Fix add() so 2+2 is 4, then retest."},
                    asst("Reading and testing.", [tc(0, "bash", {"command": "cat mathlib.py && python -m pytest -q"})]),
                    tool(0, "def add(a,b):\n    return a+b+1\nF\nassert add(2,2)==4\n"),
                    asst(
                        "Fixing the off-by-one.",
                        [
                            tc(
                                1,
                                "edit",
                                {
                                    "path": "mathlib.py",
                                    "edits": [{"oldText": "    return a+b+1", "newText": "    return a+b"}],
                                },
                            )
                        ],
                    ),
                    tool(1, "Edited"),
                    asst("Re-running tests.", [tc(2, "bash", {"command": "python -m pytest -q"})]),
                    tool(2, ".\n1 passed\n"),
                    asst("Fixed add() and tests pass."),
                ],
                "fin:fix-retest",
            )
        )

    # mix in shortened existing data
    existing = load_dataset("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-short", split="train")
    idx = list(range(len(existing)))
    rng.shuffle(idx)
    for j in idx[:8000]:
        rec = existing[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-finish")
    out.mkdir(parents=True, exist_ok=True)
    ds = 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],
        }
    )
    ds.to_parquet(str(out / "train.parquet"))
    print("wrote", len(rows), "to", out)


if __name__ == "__main__":
    main()