"""C1b: compress our own 4 JSON payloads with headroom.compress() and count tokens (Headroom's counter + tiktoken o200k_base)."""
import json, sys, time, tiktoken
from headroom import compress
enc = tiktoken.get_encoding("o200k_base")
P = {"github_issues": "payloads/gh_issues.json", "pypi_info": "payloads/pypi_info.json", "synthetic_logs": "payloads/synthetic_logs.json", "wiki_search": "payloads/wiki_search.json"}
out = {}
for name, path in P.items():
    raw = open(path).read()
    msgs = [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Here is the tool output. Answer questions about it."},
            {"role": "assistant", "content": None, "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "fetch", "arguments": "{}"}}]},
            {"role": "tool", "tool_call_id": "call_1", "content": raw}]
    t0 = time.time(); res = compress(msgs, model="gpt-4o"); dt = time.time() - t0
    comp = next(m["content"] for m in res.messages if m.get("role") == "tool")
    open(f"payloads/{name}.compressed.txt", "w").write(comp if isinstance(comp, str) else json.dumps(comp))
    b, a = len(enc.encode(raw)), len(enc.encode(comp if isinstance(comp, str) else json.dumps(comp)))
    out[name] = {"tiktoken_before": b, "tiktoken_after": a, "saved_pct_tiktoken": round(100 * (1 - a / b), 1),
                 "headroom_tokens_before": getattr(res, "tokens_before", None), "headroom_tokens_after": getattr(res, "tokens_after", None),
                 "headroom_tokens_saved": getattr(res, "tokens_saved", None), "headroom_ratio": getattr(res, "compression_ratio", None), "wall_s_first_call": round(dt, 3),
                 "transforms": str(getattr(res, "transforms_applied", ""))[:300]}
    print(name, out[name])
json.dump(out, open(sys.argv[1], "w"), indent=1)