"""C3: compress() latency on a ~10K-token and a ~100K-token JSON search result (real Wikipedia search API results).
Pre-registered: 3 warm-up calls, then 50 timed calls on the same input; report p50 (time.perf_counter around compress()).
EXTRA (labelled): 50 calls where every input is different (rotated + one changed field) so no result cache can help."""
import json, time, statistics, sys, tiktoken
from headroom import compress
enc = tiktoken.get_encoding("o200k_base")
P = "payloads/"
items = json.load(open(P + "wiki_search.json"))
for q in ("java", "linux", "database"):
    items += json.load(open(P + f"wiki_{q}.raw.json"))["query"]["search"]
def ntok(s): return len(enc.encode(s))
def take(target):
    lo = []
    for it in items:
        lo.append(it)
        if ntok(json.dumps(lo)) >= target: break
    return lo
def msgs(raw):
    return [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Search Wikipedia and summarise the results."},
            {"role": "assistant", "content": None, "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "search", "arguments": "{}"}}]},
            {"role": "tool", "tool_call_id": "call_1", "content": raw}]
def pct(xs, p): xs = sorted(xs); return xs[min(len(xs) - 1, int(round(p / 100 * (len(xs) - 1))))]
out = {}
for target in (10_000, 100_000):
    lst = take(target); raw = json.dumps(lst); m = msgs(raw)
    for _ in range(3): r = compress(m, model="gpt-4o")
    ts = []
    for _ in range(50):
        t = time.perf_counter(); r = compress(m, model="gpt-4o"); ts.append((time.perf_counter() - t) * 1000)
    vs = []
    for i in range(50):
        l2 = lst[i % len(lst):] + lst[:i % len(lst)]
        l2 = [dict(x) for x in l2]; l2[0]["size"] = l2[0]["size"] + i + 1
        m2 = msgs(json.dumps(l2))
        t = time.perf_counter(); r2 = compress(m2, model="gpt-4o"); vs.append((time.perf_counter() - t) * 1000)
    out[f"{target // 1000}K"] = {"items": len(lst), "tiktoken_o200k": ntok(raw), "headroom_tokens_before": r.tokens_before, "headroom_tokens_after": r.tokens_after,
        "saved_pct": round(100 * (r.tokens_before - r.tokens_after) / r.tokens_before, 1), "transforms": list(r.transforms_applied),
        "same_input_ms": {"p50": round(statistics.median(ts), 3), "p95": round(pct(ts, 95), 3), "min": round(min(ts), 3), "max": round(max(ts), 3)},
        "EXTRA_fresh_input_ms": {"p50": round(statistics.median(vs), 3), "p95": round(pct(vs, 95), 3), "min": round(min(vs), 3), "max": round(max(vs), 3)}}
    print(target, json.dumps(out[f"{target // 1000}K"]), flush=True)
import platform, os
out["machine"] = {"cpu_count": os.cpu_count(), "python": platform.python_version(), "headroom": __import__("headroom").__version__}
json.dump(out, open(sys.argv[1], "w"), indent=1)