"""Three query modes plus hybrid with alpha. Queries are embedded with the 'search_query: ' prefix."""
import sys
import weaviate
from weaviate.classes.query import HybridFusion, MetadataQuery

sys.path.insert(0, "../data")
from ingest import embed_query  # noqa: E402  (adds the "search_query: " prefix)

PROPS = ["heading", "text"]  # BM25 searches these two properties, not url or chunk_id
RET = ["url", "heading"]


def bm25(docs, q, k=20):
    return docs.query.bm25(query=q, query_properties=PROPS, limit=k, return_properties=RET,
                           return_metadata=MetadataQuery(score=True)).objects


def vector(docs, q, k=20):
    return docs.query.near_vector(near_vector=embed_query(q), limit=k, return_properties=RET,
                                  return_metadata=MetadataQuery(distance=True)).objects


def hybrid(docs, q, alpha=0.5, fusion=HybridFusion.RELATIVE_SCORE, k=20):
    return docs.query.hybrid(query=q, vector=embed_query(q), alpha=alpha, fusion_type=fusion,
                             query_properties=PROPS, limit=k, return_properties=RET,
                             return_metadata=MetadataQuery(score=True)).objects


def show(label, objs, n=3):
    print(label)
    for o in objs[:n]:
        m = o.metadata
        s = m.score if m.score is not None else (1 - m.distance if m.distance is not None else None)
        print(f"  {s:.3f}  {o.properties['url'].replace('https://webcrawlerapi.com', '')}  |  {o.properties['heading']}")


if __name__ == "__main__":
    q = " ".join(sys.argv[1:]) or "name_not_resolved"
    with weaviate.connect_to_local(port=8080, grpc_port=50051) as client:
        docs = client.collections.get("DocsChunk")
        print(f'query: "{q}"')
        show("bm25", bm25(docs, q))
        show("near_vector (score = 1 - cosine distance)", vector(docs, q))
        show("hybrid alpha=0.5, relative score fusion", hybrid(docs, q, 0.5))