"""C4: SAVE, restart from the snapshot, compare keys and KNN results before/after. Usage: c4_persist.py before|after out.json"""
import redis, numpy as np, json, sys, time, hashlib
r = redis.Redis(port=6380)
for _ in range(120):
    try: r.ping(); r.dbsize(); break
    except Exception: time.sleep(0.5)
if sys.argv[1] == "before":
    for i in range(1000): r.set(f"kv:{i}", f"value-{i}")
rng = np.random.default_rng(42); rng.standard_normal((10000, 64)); QV = rng.standard_normal((100, 64)).astype(np.float32)
def knn(idx):
    for _ in range(120):
        try:
            res = r.execute_command("FT.SEARCH", idx, "*=>[KNN 10 @vec $q]", "PARAMS", "2", "q", QV[0].tobytes(), "NOCONTENT", "DIALECT", "2"); break
        except redis.exceptions.BusyLoadingError: time.sleep(0.5)
    out = []
    for qi in range(20):
        res = r.execute_command("FT.SEARCH", idx, "*=>[KNN 10 @vec $q]", "PARAMS", "2", "q", QV[qi].tobytes(), "NOCONTENT", "DIALECT", "2")
        out.append([x[b"id"].decode() for x in res[b"results"]] if isinstance(res, dict) else [x.decode() for x in res[1:]])
    return out
state = {"dbsize": r.dbsize(), "kv_hash": hashlib.sha256("".join((r.get(f"kv:{i}") or b"<none>").decode() for i in range(1000)).encode()).hexdigest(),
         "vec_sample_hash": hashlib.sha256(b"".join(r.hget(f"hnsw:{i}", "vec") or b"" for i in range(0, 10000, 97))).hexdigest(),
         "ft_list": sorted(x.decode() for x in r.execute_command("FT._LIST")), "knn_flat": knn("idx_flat"), "knn_hnsw": knn("idx_hnsw"),
         "text_apple_total": None}
try:
    res = r.execute_command("FT.SEARCH", "idx_text", "@body:apple", "NOCONTENT", "LIMIT", "0", "0", "DIALECT", "2"); state["text_apple_total"] = res[b"total_results"] if isinstance(res, dict) else res[0]
except Exception as e: state["text_apple_total"] = str(e)
if sys.argv[1] == "before": state["save"] = str(r.save())
json.dump(state, open(sys.argv[2], "w"), indent=1); print({k: (v if not k.startswith("knn") else f"{len(v)} queries") for k, v in state.items()})