"""C2b: same answers? A local LLM (Qwen2.5-1.5B-Instruct Q4_K_M via Ollama 0.40.6, CPU, temperature 0, seed 1)
answers 5 code-checked questions from the ORIGINAL slice and from Headroom's COMPRESSED slice.
Slices are used because the full payloads do not fit the model's 32K context on CPU in reasonable time.
github_issues / pypi_info are skipped: Headroom returned them byte-identical (0% compression), so the model input is identical."""
import json, sys, time, urllib.request, re
from headroom import compress
P = "payloads/"
MODEL = sys.argv[2] if len(sys.argv) > 2 else "qwen15-local"
MODEL_FILE = {"qwen15-local": "qwen2.5-1.5b-instruct-q4_k_m.gguf sha256 6a1a2eb6d15622bf3c96857206351ba97e1af16c30d7a74ee38970e434e9407e", "qwen3b-local": "qwen2.5-3b-instruct-q4_k_m.gguf sha256 626b4a6678b86442240e33df819e00132d3ba7dddfe1cdc4fbb18e0a9615c62d"}[MODEL]
lg = json.load(open(P + "synthetic_logs.json"))[1250:1430]
wk = json.load(open(P + "wiki_search.json"))[:100]
def comp(raw):
    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}]
    r = compress(msgs, model="gpt-4o")
    c = r.messages[-1]["content"]
    return (json.loads(c) if c.startswith('"') else c), r
fatal = next(x for x in lg if x["level"] == "FATAL")
Q = {
 "synthetic_logs_slice": (lg, [
   ("What is the request_id of the FATAL entry?", fatal["request_id"]),
   ("Which service logged the FATAL entry?", fatal["service"]),
   ("What is the latency_ms of request %s?" % lg[10]["request_id"], str(lg[10]["latency_ms"])),
   ("What is the level of request %s?" % lg[150]["request_id"], lg[150]["level"]),
   ("Which service handled request %s?" % lg[90]["request_id"], lg[90]["service"])]),
 "wiki_search_slice": (wk, [
   ("What is the pageid of the article titled '%s'?" % wk[5]["title"], str(wk[5]["pageid"])),
   ("What is the wordcount of the article titled '%s'?" % wk[20]["title"], str(wk[20]["wordcount"])),
   ("What is the title of the article with pageid %d?" % wk[40]["pageid"], wk[40]["title"]),
   ("What is the size of the article titled '%s'?" % wk[70]["title"], str(wk[70]["size"])),
   ("What is the title of the very first result?", wk[0]["title"])]),
}
def ask(ctx, q):
    prompt = "Tool output:\n" + ctx + "\n\nQuestion: " + q + "\nAnswer with only the value, nothing else."
    body = {"model": MODEL, "prompt": prompt, "stream": False, "options": {"temperature": 0, "seed": 1, "num_ctx": 32768}}
    t = time.time()
    o = json.load(urllib.request.urlopen(urllib.request.Request("http://127.0.0.1:11434/api/generate", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}), timeout=1800))
    return o["response"].strip(), o.get("prompt_eval_count"), round(time.time() - t, 1)
def norm(s): return re.sub(r"[\s'\"`.]+", "", s).lower()
out = {"model": MODEL_FILE, "slices": {}}
for name, (rows, qs) in Q.items():
    raw = json.dumps(rows, indent=1)
    c, r = comp(raw)
    rec = {"tokens_before": r.tokens_before, "tokens_after": r.tokens_after, "transforms": list(r.transforms_applied), "ccr_markers": c.count("<<ccr:"), "questions": []}
    for label, ctx in (("original", raw), ("compressed", c)):
        for q, a in qs:
            ans, pt, s = ask(ctx, q)
            ok = norm(a) == norm(ans) or (norm(a) in norm(ans) and len(norm(ans)) <= len(norm(a)) + 12)
            rec["questions"].append({"input": label, "q": q, "truth": a, "model": ans[:200], "correct": ok, "prompt_tokens": pt, "seconds": s})
            print(name, label, ok, repr(a), repr(ans[:80]), pt, s, flush=True)
    for label in ("original", "compressed"):
        rec[label + "_correct"] = sum(1 for x in rec["questions"] if x["input"] == label and x["correct"])
    rec["same_answer_both"] = sum(1 for i in range(len(qs)) if norm(rec["questions"][i]["model"]) == norm(rec["questions"][i + len(qs)]["model"]))
    out["slices"][name] = rec
    json.dump(out, open(sys.argv[1], "w"), indent=1, ensure_ascii=False)
tot = lambda k: sum(v[k] for v in out["slices"].values())
out["summary"] = {"questions_per_input": 10, "original_correct": tot("original_correct"), "compressed_correct": tot("compressed_correct"), "identical_answers": tot("same_answer_both")}
json.dump(out, open(sys.argv[1], "w"), indent=1, ensure_ascii=False)
print(json.dumps(out["summary"]))