"""Are the two slices given to the LLM in C2b lossless after Headroom? Decode Headroom's table, resolve CCR markers, compare every row."""
import json, csv, io, re, sys
from headroom import compress
from headroom.cache.compression_store import get_compression_store
store = get_compression_store(); CCR = re.compile(r"<<ccr:([0-9a-f]+)[^>]*>>")
lg = json.load(open("payloads/synthetic_logs.json"))[1250:1430]; wk = json.load(open("payloads/wiki_search.json"))[:100]
out = {}
for name, rows in (("synthetic_logs_slice", lg), ("wiki_search_slice", wk)):
    raw = json.dumps(rows, indent=1)
    m = [{"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}]
    c = compress(m, model="gpt-4o").messages[-1]["content"]; c = json.loads(c) if c.startswith('"') else c
    marks = {}
    def ph(x): k = "@@%d@@" % len(marks); marks[k] = x.group(1); return k
    t = CCR.sub(ph, c); h = re.match(r"\[(\d+)\]\{([^}]*)\}\n", t); spec = [f.split(":") for f in h.group(2).split(",")]
    dec = []
    for r in csv.reader(io.StringIO(t[h.end():])):
        if not r: continue
        d = {}
        for (k, ty), v in zip(spec, r):
            if v in marks: v = store.retrieve(marks[v]).original_content
            d[k] = int(v) if ty == "int" else v
        dec.append(d)
    ident = sum(1 for a, b in zip(rows, dec) if {k: a[k] for k in b} == b and set(a) == set(b))
    out[name] = {"rows": len(rows), "decoded": len(dec), "identical_after_ccr": ident, "ccr_markers": len(marks), "first_300_chars_of_compressed": c[:300]}
    print(name, out[name]["rows"], out[name]["decoded"], ident, len(marks))
json.dump(out, open(sys.argv[1], "w"), indent=1, ensure_ascii=False)