#!/usr/bin/env python3
"""Convert our own solved R2E rollouts to pi-format SFT. Not eval-suite items."""

from __future__ import annotations

import json
from pathlib import Path

from datasets import Dataset

OUT = Path("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-solved2")
ROOTS = [
    Path("/mnt/pvc/users/simon/agentptb/runs/grok/workspace/outputs/rl-r2e4/run_default/rollouts"),
    Path("/mnt/pvc/users/simon/agentptb/runs/grok/workspace/outputs/rl-r2e7/run_default/rollouts"),
    Path("/mnt/pvc/users/simon/agentptb/runs/grok/workspace/outputs/rl-r2e8/run_default/rollouts"),
]
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": "read",
                "description": "Read a file",
                "parameters": {
                    "type": "object",
                    "properties": {"path": {"type": "string"}},
                    "required": ["path"],
                },
            },
        },
        {
            "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"],
                },
            },
        },
    ]
)


def score(tr, key):
    v = (tr.get("rewards") or {}).get(key)
    return float((v.get("score") if isinstance(v, dict) else v) or 0)


def norm_tc(tc: dict) -> dict:
    if "function" in tc:
        return {
            "id": tc.get("id") or "call_0",
            "type": "function",
            "function": {
                "name": (tc.get("function") or {}).get("name") or "bash",
                "arguments": (tc.get("function") or {}).get("arguments") or "{}",
            },
        }
    args = tc.get("arguments")
    if not isinstance(args, str):
        args = json.dumps(args or {})
    return {
        "id": tc.get("id") or "call_0",
        "type": "function",
        "function": {"name": tc.get("name") or "bash", "arguments": args},
    }


def has_edit(msgs) -> bool:
    for m in msgs:
        if m.get("role") != "assistant":
            continue
        for tc in m.get("tool_calls") or []:
            fn = (tc.get("function") or {}).get("name") or tc.get("name")
            if fn in {"edit", "write"}:
                return True
            args = (tc.get("function") or {}).get("arguments") or tc.get("arguments") or ""
            if isinstance(args, dict):
                args = json.dumps(args)
            if any(tok in str(args) for tok in ("sed -i", "tee ", "cat >", "git apply", "printf ")):
                return True
    return False


def to_msgs(tr) -> list[dict] | None:
    nodes = tr.get("nodes") or []
    if not (16 <= len(nodes) <= 200):
        return None
    msgs = []
    for n in nodes:
        m = n.get("message") or {}
        role = m.get("role")
        if role not in {"system", "user", "assistant", "tool"}:
            continue
        out = {"role": role}
        content = m.get("content")
        if isinstance(content, list):
            parts = []
            for p in content:
                if isinstance(p, dict) and p.get("type") == "text":
                    parts.append(p.get("text") or "")
                elif isinstance(p, str):
                    parts.append(p)
            content = "\n".join(parts)
        if content is None:
            content = ""
        if not isinstance(content, str):
            content = str(content)
        out["content"] = content
        if role == "assistant" and m.get("tool_calls"):
            out["tool_calls"] = [norm_tc(tc) for tc in m["tool_calls"]]
        if role == "tool":
            if m.get("tool_call_id"):
                out["tool_call_id"] = m["tool_call_id"]
            elif m.get("name"):
                out["name"] = m["name"]
        msgs.append(out)
    if len(msgs) < 8:
        return None
    return msgs


def main():
    rows = []
    seen = set()
    for root in ROOTS:
        if not root.exists():
            continue
        for f in root.rglob("traces.jsonl"):
            for line in f.open():
                rec = json.loads(line)
                traces = rec.get("traces") or [rec]
                for tr in traces:
                    if not isinstance(tr, dict):
                        continue
                    if score(tr, "solved") < 1:
                        continue
                    msgs = to_msgs(tr)
                    if not msgs:
                        continue
                    key = json.dumps([m.get("content", "")[:80] for m in msgs[:6]], sort_keys=True)
                    if key in seen:
                        continue
                    seen.add(key)
                    rows.append({"messages": msgs, "tools": TOOLS, "source": "r2e-solved"})
    print("rows", len(rows))
    OUT.mkdir(parents=True, exist_ok=True)
    Dataset.from_list(rows).to_parquet(str(OUT / "train.parquet"))
    print("wrote", OUT / "train.parquet")


if __name__ == "__main__":
    main()