#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2023/6/11 21:08
@Author  : hezhaozhao
@File    : test_qdrant_store.py
"""
import random

from qdrant_client.models import (
    Distance,
    FieldCondition,
    Filter,
    PointStruct,
    Range,
    VectorParams,
)

from metagpt.document_store.qdrant_store import QdrantConnection, QdrantStore

seed_value = 42
random.seed(seed_value)

vectors = [[random.random() for _ in range(2)] for _ in range(10)]

points = [
    PointStruct(id=idx, vector=vector, payload={"color": "red", "rand_number": idx % 10})
    for idx, vector in enumerate(vectors)
]


def test_qdrant_store():
    qdrant_connection = QdrantConnection(memory=True)
    vectors_config = VectorParams(size=2, distance=Distance.COSINE)
    qdrant_store = QdrantStore(qdrant_connection)
    qdrant_store.create_collection("Book", vectors_config, force_recreate=True)
    assert qdrant_store.has_collection("Book") is True
    qdrant_store.delete_collection("Book")
    assert qdrant_store.has_collection("Book") is False
    qdrant_store.create_collection("Book", vectors_config)
    assert qdrant_store.has_collection("Book") is True
    qdrant_store.add("Book", points)
    results = qdrant_store.search("Book", query=[1.0, 1.0])
    assert results[0]["id"] == 2
    assert results[0]["score"] == 0.999106722578389
    assert results[1]["id"] == 7
    assert results[1]["score"] == 0.9961650411397226
    results = qdrant_store.search("Book", query=[1.0, 1.0], return_vector=True)
    assert results[0]["id"] == 2
    assert results[0]["score"] == 0.999106722578389
    assert results[0]["vector"] == [0.7363563179969788, 0.6765939593315125]
    assert results[1]["id"] == 7
    assert results[1]["score"] == 0.9961650411397226
    assert results[1]["vector"] == [0.7662628889083862, 0.6425272226333618]
    results = qdrant_store.search(
        "Book",
        query=[1.0, 1.0],
        query_filter=Filter(must=[FieldCondition(key="rand_number", range=Range(gte=8))]),
    )
    assert results[0]["id"] == 8
    assert results[0]["score"] == 0.9100373450784073
    assert results[1]["id"] == 9
    assert results[1]["score"] == 0.7127610621127889
    results = qdrant_store.search(
        "Book",
        query=[1.0, 1.0],
        query_filter=Filter(must=[FieldCondition(key="rand_number", range=Range(gte=8))]),
        return_vector=True,
    )
    assert results[0]["vector"] == [0.35037919878959656, 0.9366079568862915]
    assert results[1]["vector"] == [0.9999677538871765, 0.00802854634821415]
