"""C2: vector (FLAT + HNSW) and full-text/tag/numeric search on KiviDB vs exact ground truth computed in Python."""
import redis, numpy as np, json, time, random, sys
r = redis.Redis(port=6380); r.flushall()
def P(res):
    """redis-py 8 turns FT.SEARCH replies into a dict (client-side callback; the server sends a normal RESP2 array, checked with a raw socket). Return (total, [(id, fields)])."""
    if isinstance(res, dict):
        return res[b"total_results"], [(x[b"id"].decode(), {k.decode(): v.decode() for k, v in x.get(b"extra_attributes", {}).items()}) for x in res[b"results"]]
    total, rest = res[0], res[1:]
    if rest and isinstance(rest[1] if len(rest) > 1 else None, list):
        return total, [(rest[i].decode(), {rest[i+1][j].decode(): rest[i+1][j+1].decode() for j in range(0, len(rest[i+1]), 2)}) for i in range(0, len(rest), 2)]
    return total, [(x.decode(), {}) for x in rest]
N, D, Q, K = 10000, 64, 100, 10
rng = np.random.default_rng(42); X = rng.standard_normal((N, D)).astype(np.float32); QV = rng.standard_normal((Q, D)).astype(np.float32)
Xn = X / np.linalg.norm(X, axis=1, keepdims=True); Qn = QV / np.linalg.norm(QV, axis=1, keepdims=True)
truth = np.argsort(-(Qn @ Xn.T), axis=1)[:, :K]
out = {}
for algo, extra in (("FLAT", []), ("HNSW", ["M", "16", "EF_CONSTRUCTION", "200"])):
    pre = f"{algo.lower()}:"; idx = f"idx_{algo.lower()}"
    nargs = 6 + len(extra)
    out[f"{algo}_create"] = str(r.execute_command("FT.CREATE", idx, "ON", "HASH", "PREFIX", "1", pre, "SCHEMA", "vec", "VECTOR", algo, str(nargs), "TYPE", "FLOAT32", "DIM", str(D), "DISTANCE_METRIC", "COSINE", *extra))
    t0 = time.time(); p = r.pipeline(transaction=False)
    for i in range(N):
        p.hset(f"{pre}{i}", mapping={"vec": X[i].tobytes()})
        if i % 1000 == 999: p.execute()
    p.execute(); ingest = time.time() - t0
    time.sleep(1)
    hits = 0; lat = []; bad = 0
    for qi in range(Q):
        t1 = time.time()
        res = r.execute_command("FT.SEARCH", idx, f"*=>[KNN {K} @vec $q AS score]", "PARAMS", "2", "q", QV[qi].tobytes(), "SORTBY", "score", "RETURN", "1", "score", "LIMIT", "0", str(K), "DIALECT", "2")
        lat.append((time.time() - t1) * 1000)
        ids = [int(i.split(":")[1]) for i, _ in P(res)[1]]
        if len(ids) != K: bad += 1
        hits += len(set(ids) & set(truth[qi].tolist()))
    out[algo] = {"recall_at_10": round(hits / (Q * K), 4), "queries_with_wrong_count": bad, "ingest_s": round(ingest, 2), "query_ms_p50": round(float(np.median(lat)), 2)}
    # check scores against our own cosine distance for query 0
    res = r.execute_command("FT.SEARCH", idx, f"*=>[KNN 3 @vec $q AS score]", "PARAMS", "2", "q", QV[0].tobytes(), "SORTBY", "score", "RETURN", "1", "score", "DIALECT", "2")
    out[algo]["q0_top3"] = [(i, float(f["score"])) for i, f in P(res)[1]]
    out[algo]["q0_top3_truth"] = [(f"{pre}{i}", round(float(1 - Qn[0] @ Xn[i]), 6)) for i in truth[0][:3]]
# full text / tag / numeric
words = ["apple", "banana", "cherry", "delta", "echo", "fox", "golf", "hotel"]; colors = ["red", "green", "blue"]
random.seed(7); docs = {}
for i in range(200):
    w = random.sample(words, 3); docs[i] = {"body": " ".join(w + ["filler", "text"]), "color": random.choice(colors), "price": random.randint(1, 50)}
    r.hset(f"doc:{i}", mapping=docs[i])
out["text_create"] = str(r.execute_command("FT.CREATE", "idx_text", "ON", "HASH", "PREFIX", "1", "doc:", "SCHEMA", "body", "TEXT", "color", "TAG", "price", "NUMERIC"))
time.sleep(1)
def ft(q):
    res = r.execute_command("FT.SEARCH", "idx_text", q, "NOCONTENT", "LIMIT", "0", "1000", "DIALECT", "2")
    total, items = P(res); return total, sorted(int(i.split(":")[1]) for i, _ in items)
cases = {"@body:apple": lambda d: "apple" in d["body"].split(),
         "@body:(apple banana)": lambda d: {"apple", "banana"} <= set(d["body"].split()),
         "@body:(apple|cherry)": lambda d: bool({"apple", "cherry"} & set(d["body"].split())),
         "@body:apple -@body:banana": lambda d: "apple" in d["body"].split() and "banana" not in d["body"].split(),
         "@color:{red}": lambda d: d["color"] == "red",
         "@price:[10 20]": lambda d: 10 <= d["price"] <= 20,
         "@body:fox @color:{blue} @price:[30 +inf]": lambda d: "fox" in d["body"].split() and d["color"] == "blue" and d["price"] >= 30}
tx = {}
for q, f in cases.items():
    want = sorted(i for i, d in docs.items() if f(d))
    try: total, got = ft(q); tx[q] = {"expected": len(want), "got": len(got), "total_reported": total, "exact": got == want}
    except Exception as e: tx[q] = {"error": str(e)}
out["text"] = tx
# BM25 ranking sanity: doc containing the term twice should rank above docs with it once
r.hset("doc:900", mapping={"body": "golf golf golf filler", "color": "red", "price": 5}); time.sleep(0.5)
res = r.execute_command("FT.SEARCH", "idx_text", "@body:golf", "NOCONTENT", "LIMIT", "0", "3", "DIALECT", "2")
out["bm25_top_for_golf"] = [i for i, _ in P(res)[1]]
json.dump(out, open(sys.argv[1], "w"), indent=1); print(json.dumps(out, indent=1))