#!/usr/bin/env python3
"""Compact one-turn pi traces for LoRA on top of sft-v3. Generic skills, not eval items."""
from __future__ import annotations

import json
import random
from pathlib import Path

from datasets import Dataset, load_dataset

SYS = (
    "You are a coding agent. Prefer tools over talk. "
    "Chain work in one bash call. Never loop ls/cd/pwd. "
    "After the required output exists, one short sentence and 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": "edit",
                "description": "Edit a file",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}, "edits": {"type": "array"}},
                    "required": ["path", "edits"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "read",
                "description": "Read a file",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}},
                    "required": ["path"],
                },
            },
        },
    ]
)


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


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


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

    # extract archive to the destination named in the prompt (not a nested vendor dir)
    for i in range(200):
        rows.append(
            pack(
                [
                    {
                        "role": "user",
                        "content": (
                            f"A tarball is at /data/pkg{i}.tar.gz. Extract it to /workspace/pkg, "
                            "run `make -C /workspace/pkg && make -C /workspace/pkg install PREFIX=/output`, then stop."
                        ),
                    },
                    {
                        "role": "assistant",
                        "content": "Extracting to /workspace/pkg and building.",
                        "tool_calls": [
                            tc(
                                0,
                                "bash",
                                {
                                    "command": (
                                        f"mkdir -p /workspace/pkg /output && "
                                        f"tar -xzf /data/pkg{i}.tar.gz -C /workspace/pkg --strip-components=1 && "
                                        "make -C /workspace/pkg && make -C /workspace/pkg install PREFIX=/output"
                                    )
                                },
                            )
                        ],
                    },
                    {"role": "tool", "tool_call_id": "c0", "content": "install done\n"},
                    {"role": "assistant", "content": "Built and installed to /output."},
                ],
                "lora:extract",
            )
        )

    # compact write via bash heredoc (fits in ~120 tokens)
    for i in range(200):
        rows.append(
            pack(
                [
                    {
                        "role": "user",
                        "content": (
                            "Write /output/strip.py that reads stdin and prints lines that do not contain "
                            "the substring TODO. Then stop."
                        ),
                    },
                    {
                        "role": "assistant",
                        "content": "Writing the filter script.",
                        "tool_calls": [
                            tc(
                                0,
                                "bash",
                                {
                                    "command": (
                                        "cat > /output/strip.py << 'PY'\n"
                                        "import sys\n"
                                        "for line in sys.stdin:\n"
                                        "    if 'TODO' not in line:\n"
                                        "        sys.stdout.write(line)\n"
                                        "PY\nchmod +x /output/strip.py"
                                    )
                                },
                            )
                        ],
                    },
                    {"role": "tool", "tool_call_id": "c0", "content": ""},
                    {"role": "assistant", "content": "Wrote /output/strip.py."},
                ],
                "lora:heredoc",
            )
        )

    # read then edit immediately (no second ls)
    for i in range(180):
        rows.append(
            pack(
                [
                    {"role": "user", "content": "off_by_one in /workspace/add.py: add(2,2) must be 4. Fix it and stop."},
                    {
                        "role": "assistant",
                        "content": "Reading then patching.",
                        "tool_calls": [tc(0, "bash", {"command": "cat /workspace/add.py"})],
                    },
                    {"role": "tool", "tool_call_id": "c0", "content": "def add(a, b):\n    return a + b + 1\n"},
                    {
                        "role": "assistant",
                        "content": "Patching the return.",
                        "tool_calls": [
                            tc(
                                1,
                                "edit",
                                {
                                    "path": "/workspace/add.py",
                                    "edits": [{"oldText": "    return a + b + 1\n", "newText": "    return a + b\n"}],
                                },
                            )
                        ],
                    },
                    {"role": "tool", "tool_call_id": "c1", "content": "Edited add.py"},
                    {"role": "assistant", "content": "add(2,2) is 4."},
                ],
                "lora:edit",
            )
        )

    # git restore previous commit (generic)
    for i in range(120):
        rows.append(
            pack(
                [
                    {
                        "role": "user",
                        "content": "The latest commit broke page.html. Restore that file from HEAD~1 and stop.",
                    },
                    {
                        "role": "assistant",
                        "content": "Restoring page.html from the previous commit.",
                        "tool_calls": [tc(0, "bash", {"command": "git checkout HEAD~1 -- page.html && git status --short"})],
                    },
                    {"role": "tool", "tool_call_id": "c0", "content": "M  page.html\n"},
                    {"role": "assistant", "content": "Restored page.html from HEAD~1."},
                ],
                "lora:git",
            )
        )

    # mix original finish set so we keep stop-after-verify
    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[:3500]:
            rec = ds[int(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-lora")
    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()