"""Hit table for Typesense: keyword, vector-only, hybrid at the default alpha (0.3) and at 0.5, plus an alpha sweep."""
import json
from common import EXACT, FETCH, TOP, hit_rank, pages, qvec, questions, ts

QS = questions() + [EXACT]


def search(q, alpha=None, vec=False, extra=None, collection="docs"):
    s = {"collection": collection, "per_page": FETCH, "exclude_fields": "embedding", "include_fields": "url,heading"}
    if vec:
        s.update(q="*", query_by="heading,text")  # vector-only: wildcard query plus nearest neighbours
        s["vector_query"] = f"embedding:({json.dumps(qvec(q))}, k:{FETCH})"
    else:
        s.update(q=q, query_by="heading,text")
        if alpha is not None:
            a = "" if alpha == "default" else f", alpha:{alpha}"
            s["vector_query"] = f"embedding:({json.dumps(qvec(q))}, k:{FETCH}{a})"
    s.update(extra or {})
    return ts("POST", "/multi_search", {"searches": [s]})["results"][0]


def urls(res):
    return [h["document"]["url"] for h in res["hits"]]


if __name__ == "__main__":
    print("keyword settings: defaults (num_typos 2, prefix true, drop_tokens_threshold 1, typo_tokens_threshold 1); query_by heading,text; fetch", FETCH, "chunks, top", TOP, "pages")
    print("modes: keyword = no vector_query; vector = q=* + vector_query; hybrid default = alpha unset (0.3 per docs); hybrid 0.5 = alpha 0.5")
    modes = {"keyword": lambda q: search(q), "vector": lambda q: search(q, vec=True),
             "hyb_def": lambda q: search(q, "default"), "hyb_0.5": lambda q: search(q, 0.5)}
    print(f"{'question':<70} " + " ".join(f"{m:>8}" for m in modes))
    tot6 = {m: 0 for m in modes}
    top5 = {}
    for i, q in enumerate(QS):
        row = []
        for m, fn in modes.items():
            res = fn(q["question"])
            r = hit_rank(urls(res), q["expected_url_contains"])
            tot6[m] += r is not None and i < 6
            row.append(f"{r or '-':>8}")
            if m == "vector":
                top5[q["question"]] = pages(urls(res))[:TOP]
        print(f"{q['question'][:68]:<70} " + " ".join(row))
    print("hits x/6 (six questions):", {m: f"{v}/6" for m, v in tot6.items()})
    print("429 rank:", {m: hit_rank(urls(fn('429')), EXACT['expected_url_contains']) for m, fn in modes.items()})
    json.dump(top5, open("ts_vector_top5.json", "w"), indent=1)

    print("\n--- alpha sweep: rank of expected page (- = miss), then x/6 hits and 429 rank ---")
    for a in (0.0, 0.25, 0.5, 0.75, 1.0):
        ranks = [hit_rank(urls(search(q["question"], a)), q["expected_url_contains"]) for q in QS]
        print(f"alpha {a:<5}", " ".join(f"{r or '-':>2}" for r in ranks[:6]), f"| {sum(r is not None for r in ranks[:6])}/6 | 429: {ranks[6] or '-'}")
    res = search(QS[0]["question"], "default")
    print("\nfound (keyword matches for hybrid, question 1):", res["found"], "| search_time_ms", res["search_time_ms"])

    print("\n--- side check, NOT in the table: keyword mode with drop_tokens_threshold:50 (default is 1) ---")
    ranks = [hit_rank(urls(search(q["question"], extra={"drop_tokens_threshold": 50})), q["expected_url_contains"]) for q in QS]
    print(" ".join(f"{r or '-':>2}" for r in ranks[:6]), f"| {sum(r is not None for r in ranks[:6])}/6 | 429: {ranks[6] or '-'}")