#!/usr/bin/env python3
"""Offline smoke test: does the served checkpoint drive the pi loop correctly?

No sandbox needed. Replays prefixes of held-out converted trajectories against the
served model and reports, over N samples:
  * fraction of turns that emit a syntactically valid tool call
  * fraction that use a real pi tool name (read/bash/edit/write)
  * completion length distribution and finish_reason (i.e. does it stop?)
  * fraction that hallucinate a tool response / continue past its own turn

Usage: check_policy.py --base-url http://localhost:8000/v1 --model policy -n 40
"""

import argparse
import collections
import glob
import json
import random
import re
import sys
import urllib.request

PI_TOOLS = {"read", "bash", "edit", "write"}
LEAK = re.compile(r"<\|im_start\|>|<\|im_end\|>|<tool_response>")


def post(url, payload, timeout=300):
    req = urllib.request.Request(
        url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-url", default="http://localhost:8000/v1")
    ap.add_argument("--model", default="policy")
    ap.add_argument("--data", default="/mnt/pvc/users/simon/agentptb/runs/a-opus-max/workspace/data/sft_v1")
    ap.add_argument("-n", type=int, default=40)
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--max-tokens", type=int, default=2048)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    random.seed(args.seed)

    import pyarrow.parquet as pq

    f = sorted(glob.glob(args.data + "/*.parquet"))[-1]
    rows = pq.ParquetFile(f).read_row_group(0).to_pylist()
    random.shuffle(rows)

    stats = collections.Counter()
    lens = []
    for row in rows[: args.n]:
        msgs = row["messages"]
        tools = json.loads(row["tool_defs"])
        # cut at a random assistant turn: the model must produce that turn
        idxs = [i for i, m in enumerate(msgs) if m["role"] == "assistant"]
        if len(idxs) < 2:
            continue
        cut = random.choice(idxs[:-1])
        prefix = []
        for m in msgs[:cut]:
            mm = {k: v for k, v in m.items() if k != "reasoning_content" and v not in (None, "")}
            if m["role"] == "assistant" and m.get("reasoning_content") and not m.get("content"):
                mm["content"] = m["reasoning_content"]
            prefix.append(mm)
        if not prefix or prefix[-1]["role"] == "assistant":
            continue
        try:
            out = post(
                args.base_url + "/chat/completions",
                {
                    "model": args.model,
                    "messages": prefix,
                    "tools": tools,
                    "temperature": args.temperature,
                    "max_tokens": args.max_tokens,
                },
            )
        except Exception as e:
            stats["http_error"] += 1
            print("ERR", repr(e)[:200], file=sys.stderr)
            continue
        ch = out["choices"][0]
        msg = ch["message"]
        stats["n"] += 1
        stats["finish_" + str(ch.get("finish_reason"))] += 1
        lens.append(out["usage"]["completion_tokens"])
        tcs = msg.get("tool_calls") or []
        if tcs:
            stats["has_tool_call"] += 1
            for tc in tcs:
                name = (tc.get("function") or {}).get("name")
                stats["tool_" + str(name)] += 1
                if name in PI_TOOLS:
                    stats["valid_tool_name"] += 1
                try:
                    json.loads((tc.get("function") or {}).get("arguments") or "{}")
                    stats["valid_args"] += 1
                except Exception:
                    stats["bad_args"] += 1
        else:
            stats["text_only"] += 1
        text = msg.get("content") or ""
        if LEAK.search(text):
            stats["leaked_special_tokens"] += 1

    lens.sort()
    n = max(1, stats["n"])
    print(json.dumps({k: v for k, v in sorted(stats.items())}, indent=1))
    if lens:
        print(
            f"completion_tokens: mean={sum(lens)//len(lens)} median={lens[len(lens)//2]} "
            f"p90={lens[int(0.9*len(lens))]} max={lens[-1]}"
        )
    print(f"tool-call rate {stats['has_tool_call']/n:.2f}  valid-name rate {stats['valid_tool_name']/max(1,stats['has_tool_call']):.2f}  leak rate {stats['leaked_special_tokens']/n:.2f}")


if __name__ == "__main__":
    main()