#!/usr/bin/env python3
"""Show that the RL rollout path and the serving path render different system prompts.

The renderer serialises whatever tool objects it is handed. verifiers hands it its own flat
`ToolSpec`, so the tool block comes out as {"name":…,"description":…,"parameters":…}; the served
chat template emits {"type":"function","function":{…}}. That is a ~68-character difference in the
system prompt on *every turn* of an RL rollout, which is why GRPO through pi scores 0 on tasks
the same harness solves ~16% of the time under the eval client.

Fix (in shared verifiers, so not done here): hand the renderer OAI-nested tools.

    python3 scripts/check_render_parity.py [ckpt] [trace.jsonl]
"""
import copy, json, sys

from transformers import AutoTokenizer
from renderers.base import create_renderer
from renderers.configs import Qwen35RendererConfig

WS = "/mnt/pvc/users/simon/agentptb/runs/a-opus-max/workspace"
ckpt = sys.argv[1] if len(sys.argv) > 1 else f"{WS}/ckpt/sft_v5/weights/step_900"
trace = sys.argv[2] if len(sys.argv) > 2 else f"{WS}/runs/verify_ws/traces.jsonl"

tk = AutoTokenizer.from_pretrained(ckpt)
rend = create_renderer(tk, Qwen35RendererConfig())
t = [json.loads(l) for l in open(trace)][0]["traces"][0]
msgs = [n["message"] for n in t["nodes"][:6]]
tools = t["tools"]
oai = [{"type": "function", "function": {"name": x["name"], "description": x.get("description", ""),
                                         "parameters": x.get("parameters")}} for x in tools]
srv_msgs = []
for m in copy.deepcopy(msgs):
    if m.get("tool_calls"):
        m["tool_calls"] = [{"id": tc.get("id"), "type": "function",
                            "function": {"name": tc.get("name"),
                                         "arguments": json.loads(tc.get("arguments") or "{}")}}
                           for tc in m["tool_calls"]]
    if isinstance(m.get("content"), list):
        m["content"] = "".join(p.get("text", "") for p in m["content"] if isinstance(p, dict))
    srv_msgs.append(m)

served = tk.apply_chat_template(srv_msgs, tools=oai, tokenize=False, add_generation_prompt=True)
flat = tk.decode(list(rend.render_ids(msgs, tools=tools, add_generation_prompt=True)))
nested = tk.decode(list(rend.render_ids(msgs, tools=oai, add_generation_prompt=True)))
print("renderer(flat ToolSpec)  == served template :", flat == served)
print("renderer(OAI-nested)     == served template :", nested == served)
print("system-prompt delta (chars):", len(served) - len(flat))