import json
import math

import chromadb

from common import HOST, PORT, COLLECTION, read_chunks, embed_query

rows, model, dim = read_chunks()
norms = [math.sqrt(sum(x * x for x in r["embedding"])) for r in rows]
print(f"vector norms: min {min(norms):.7f} max {max(norms):.7f}")

col = chromadb.HttpClient(host=HOST, port=PORT).get_collection(COLLECTION)


def top(question, n=5, **kw):
    vec = embed_query(question, model)
    res = col.query(query_embeddings=[vec], n_results=n,
                    include=["documents", "metadatas", "distances"], **kw)
    # every field is a list of lists, one inner list per query vector
    return list(zip(res["metadatas"][0], res["distances"][0], res["documents"][0]))


for q in ["How do I get notified when a crawl job finishes?", "How long are scrape results cached and how do I disable the cache?"]:
    print("\nQ:", q)
    for m, d, _ in top(q, 3):
        print(f"  {d:.3f}  {m['section']:15s} {m['url'].replace('https://webcrawlerapi.com', '')}  [{m['heading']}]")

print("\nsmoke test, 6 questions, expected URL in top 5:")
hits = []
for item in json.load(open("questions.json")):
    res = top(item["question"], 5)
    rank = next((i + 1 for i, (m, _, _) in enumerate(res) if item["expected_url_contains"] in m["url"]), None)
    hits.append(rank)
    print(f"  rank {rank}  d={res[0][1]:.3f}  {item['question']}")
print("ranks:", hits, "| hit in top 5:", sum(1 for h in hits if h), "/", len(hits))