from pymilvus import (
    MilvusClient,
    AnnSearchRequest,
    WeightedRanker
)
import os
from openai import OpenAI
from typing import List
def insert():

    docs = [
        "Artificial intelligence was founded as an academic discipline in 1956.",
        "Alan Turing was the first person to conduct substantial research in AI.",
        "Born in Maida Vale, London, Turing was raised in southern England.",
    ]

    data = [
        {"text": docs[0], "dense": [2.7242085933685303, 6.021071434020996, 0.4754035174846649, 9.358858108520508, 5.173221111297607]},
        {"text": docs[1], "dense": [8.584294319152832, 2.7640628814697266, 9.558855056762695, 2.584272861480713, 4.705013275146484]},
        {"text": docs[2], "dense": [2.5525057315826416, 3.8815805912017822, 9.343480110168457, 7.888997554779053, 4.500918388366699]},
    ]

    client = MilvusClient(
            uri="http://172.18.107.78:19530",
            token="root:Milvus"
        )
    client.use_database("JXTest")
    res = client.insert(
        collection_name="hybrid_search_collection",
        data=data
    )
    return res

def search():
    from pymilvus import AnnSearchRequest
    client = MilvusClient(
        uri="http://172.18.127.124:19530",
        token="root:Milvus"
    )
    search_param_1 = {
        "data": [[0.7425515055656433, 7.774101734161377, 0.7397570610046387, 2.429982900619507, 3.8253049850463867]],
        "anns_field": "dense",
        "param": {
            "metric_type": "L2",
            "params": {"nprobe": 10}
        },
        "limit": 2
    }
    request_1 = AnnSearchRequest(**search_param_1)

    search_param_2 = {
        "data": ['Who started AI research'],
        "anns_field": "sparse",
        "param": {
            "metric_type": "BM25",
        },
        "limit": 2
    }
    request_2 = AnnSearchRequest(**search_param_2)

    reqs = [request_1, request_2]
    return reqs

openai_embedding_api_base = os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.127.124:40697/v1")
openai_api_key = os.getenv("OPENAI_API_KEY", "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5")

embedding_client = OpenAI(
api_key=openai_api_key,
base_url=openai_embedding_api_base,
)
def get_embeddings(embedding_texts: List[str]):
    response = embedding_client.embeddings.create(
        model="bge-m3",
        input=embedding_texts,
    )
    emebdding_list = response.data[0].embedding
    return emebdding_list

def test():
    from pymilvus import AnnSearchRequest
    client = MilvusClient(
        uri="http://172.18.127.124:19530",
        token="root:Milvus"
    )
    client.use_database("JXTest")
    query_vector = get_embeddings(["焦煤谜团中央医院收治了多少名颅内动脉狭窄患者？"])
    search_param_1 = {
    "data": [query_vector],
    "anns_field": "vector",
    "param": {
        "metric_type": "COSINE",
        "params": {"nprobe": 10}
    },
    "limit": 3
}
    request_1 = AnnSearchRequest(**search_param_1)

    search_param_2 = {
        "data": ["焦煤谜团中央医院收治了多少名颅内动脉狭窄患者？"],
        "anns_field": "bm_25",
        "param": {
            "metric_type": "BM25",
        },
        "limit": 3
    }
    request_2 = AnnSearchRequest(**search_param_2)

    reqs = [request_1, request_2]
    ranker = WeightedRanker(0.5, 0.5)
    res = client.hybrid_search(
    collection_name="XianTest",
    reqs=reqs,
    ranker=ranker,
    limit=3,
    output_fields=["Header_1","Header_2","Header_3","text","resource","img_path"]
    )
    search_results = []
    search_resources = []
    for hits in res:
        for hit in hits:
            search_result_text = hit["entity"].get("text")

            search_result_resource = hit["entity"].get("resource")
            search_result_img_path = hit["entity"].get("img_path")
            search_results.append(search_result_text)
            search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path})
    client.close()
    return search_results,search_resources

if __name__ == '__main__':
    search_results,search_resources=test()
    print(search_results)
    print(search_resources)



    # insert_res = insert()
    # print(insert_res)
    # reqs = search()
    # client = MilvusClient(
    #         uri = "http://172.18.127.124:19530",
    #         # uri="http://172.18.107.78:19530",
    #         token="root:Milvus"
    #     )
    # client.use_database("JXTest")
    # print(client.list_collections())

    # from pymilvus import WeightedRanker

    # ranker = WeightedRanker(0.5, 0.5)
    # res = client.hybrid_search(
    # collection_name="hybrid_search_collection",
    # reqs=reqs,
    # ranker=ranker,
    # limit=2,
    # output_fields=["text"]
    # )
    # for hits in res:
    #     print("TopK results:")
    #     for hit in hits:
    #         print(hit)
    # Ouput :
    #{'id': 457020197413365373, 'distance': 0.3451843857765198, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.'}}
    #{'id': 457020197413365372, 'distance': 0.005594015121459961, 'entity': {'text': 'Artificial intelligence was founded as an academic discipline in 1956.'}}