import json
import os
import time

import lancedb

from common import connect, TABLE, read_chunks, to_record, embed_query, URI

rows, model, dim = read_chunks()
print("target:", URI)
db = connect()
t0 = time.time()
table = db.create_table(TABLE, data=[to_record(r) for r in rows], mode="overwrite")
print(f"loaded in {time.time() - t0:.2f}s, count: {table.count_rows()}")
assert table.count_rows() == len(rows)
print("versions:", len(table.list_versions()))

local = lancedb.connect(os.environ["LOCAL_URI"]).open_table(TABLE)
print("\nsame query, local table vs bucket table:")
same = 0
ranks = []
for item in json.load(open("questions.json")):
    v = embed_query(item["question"], model)
    a = local.search(v).distance_type("cosine").select(["url", "_distance"]).limit(5).to_list()
    b = table.search(v).distance_type("cosine").select(["url", "_distance"]).limit(5).to_list()
    same += a[0]["url"] == b[0]["url"] and abs(a[0]["_distance"] - b[0]["_distance"]) < 1e-6
    ranks.append(next((i + 1 for i, r in enumerate(b) if item["expected_url_contains"] in r["url"]), None))
    print(f"  top-1 {b[0]['url'].replace('https://webcrawlerapi.com', '')}  d={b[0]['_distance']:.3f}")
print(f"top-1 URL and distance identical to local: {same} / 6")
print("expected-page ranks on the bucket table:", ranks)

from lancedb.index import FTS
table.create_index("text", config=FTS())
hit = table.search("429", query_type="fts", fts_columns="text").select(["url", "_score"]).limit(1).to_list()[0]
print("\nFTS index on the bucket table, query '429':", hit["url"].replace("https://webcrawlerapi.com", ""))