import json
from lancedb.index import FTS
from lancedb.query import PhraseQuery
from lancedb.rerankers import RRFReranker

from common import connect, TABLE, read_chunks, embed_query

rows, model, dim = read_chunks()
table = connect().open_table(TABLE)
table.create_index("text", config=FTS(), replace=True)
print("indices:", [(i.name, i.index_type, i.columns) for i in table.list_indices()])


def short(u):
    return u.replace("https://webcrawlerapi.com", "")


def fts(q, n=3):
    return (table.search(q, query_type="fts", fts_columns="text")
            .select(["url", "heading", "text", "_score"]).limit(n).to_list())


def vec(q, n=3):
    return (table.search(embed_query(q, model)).distance_type("cosine")
            .select(["url", "heading", "text", "_distance"]).limit(n).to_list())


def hybrid(q, n=3):
    return (table.search(query_type="hybrid", vector_column_name="vector", fts_columns="text")
            .vector(embed_query(q, model)).text(q).rerank(RRFReranker())
            .limit(n).to_list())


print("\n== how is max_depth tokenized? ==")
n_exact = sum(1 for r in rows if "max_depth" in r["text"])
print("chunks containing the string max_depth:", n_exact)
for q in ["max_depth", "depth", "max"]:
    res = fts(q, 12)
    print(f"  fts {q!r:12s} returned {len(res)} rows, {sum(1 for r in res if 'max_depth' in r['text'])} contain 'max_depth'")
top_a = [r["id"] if "id" in r else r["url"] + r["heading"] for r in fts("max_depth", 5)]
top_b = [r["id"] if "id" in r else r["url"] + r["heading"] for r in fts("depth", 5)]
print("  top 5 of 'max_depth' == top 5 of 'depth':", top_a == top_b)
def phrase():
    return (table.search(PhraseQuery("max_depth", "text"), query_type="fts")
            .select(["url", "heading", "text", "_score"]).limit(12).to_list())


try:
    phrase()
except ValueError as e:
    print("  PhraseQuery on the default index fails:", str(e))
table.create_index("text", config=FTS(with_position=True), replace=True)
ph = phrase()
print("  PhraseQuery('max_depth'):", len(ph), "rows,", sum(1 for r in ph if "max_depth" in r["text"]), "contain 'max_depth'")

for kw in ["whitelist_regexp", "429"]:
    print(f"\n== keyword-style query: {kw} ==")
    for name, fn, key in [("fts", fts, "_score"), ("vector", vec, "_distance"), ("hybrid", hybrid, None)]:
        print(f"{name}:")
        for r in fn(kw, 3):
            s = f"{r[key]:.2f}" if key else f"{r.get('_relevance_score', 0):.4f}"
            print(f"  {s:>7s}  {short(r['url'])}  [{r['heading'][:40]}]  has term: {kw in r['text']}")

print("\n== natural-language question ==")
q = "How do I make a crawl only visit certain URLs?"
for name, fn, key in [("fts", fts, "_score"), ("vector", vec, "_distance"), ("hybrid", hybrid, None)]:
    print(f"{name}:")
    for r in fn(q, 3):
        s = f"{r[key]:.2f}" if key else f"{r.get('_relevance_score', 0):.4f}"
        print(f"  {s:>7s}  {short(r['url'])}  [{r['heading'][:40]}]")

print("\n== hybrid, 6 questions, expected URL in top 5 ==")
for name, fn in [("vector", vec), ("fts", fts), ("hybrid", hybrid)]:
    ranks = []
    for item in json.load(open("questions.json")):
        res = fn(item["question"], 5)
        ranks.append(next((i + 1 for i, r in enumerate(res) if item["expected_url_contains"] in r["url"]), None))
    print(f"  {name:7s} ranks {ranks}  hit {sum(1 for x in ranks if x)}/6")
print("\nhybrid result keys:", sorted(hybrid("whitelist_regexp", 1)[0].keys()))