import os
from pymilvus import (
    connections,
    utility,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    MilvusClient,
    db,
    Function,
    FunctionType,
    AnnSearchRequest,
    WeightedRanker
)

class Milvus_Database:
    def __init__(self,user=os.getenv("Milvus_USER","root"),password=os.getenv("Milvus_PASSWORD","Milvus"),uri="http://172.18.30.165:19530"):
        self.client = MilvusClient(
        uri="http://172.18.30.165:19530",
        token="root:Milvus"
    )
    def get_text_by_resource(self, collection_name, database_name,resource,text):
        """
        return List :
        """
        self.client.using_database(db_name=database_name)
        iterator = self.client.query_iterator(
            collection_name=collection_name,
            batch_size=5,
            filter=f"text like \"{text}\" and resource like \"{resource}\"",
            output_fields=["id","resource","text"],
        )

        results = []

        while True:
            result = iterator.next()
            if not result:
                iterator.close()
                break
            results += result
        return results
    def get_chunk_by_id(self, collection_name, database_name,id):
        """
        return List :
        """
        try:
            self.client.using_database(db_name=database_name)
            if not self.client.has_collection(collection_name):
                return []
            else:
                self.client.load_collection(collection_name)
            iterator = self.client.query_iterator(
                collection_name=collection_name,
                batch_size=5,
                filter=f"id=={id}",
                output_fields=["id","text","resource"],
            )

            results = []

            while True:
                result = iterator.next()
                if not result:
                    iterator.close()
                    break
                results += result
            return results
        finally:
                # 3. 无论成功还是失败，最终释放集合
                try:
                    self.client.release_collection(collection_name)
                except Exception as e:
                    print(f"释放集合失败: {e}")

def get_chunk(collection_name,database_name,resource,text,id):
    milvus_db=Milvus_Database()
    # results = milvus_db.get_text_by_resource(collection_name=collection_name,database_name=database_name,resource=resource,text=text)
    results = milvus_db.get_chunk_by_id(collection_name=collection_name,database_name=database_name,id=id)
    return {"code": 200, "message": "success", "data": results}

if __name__ == "__main__":
    text = "你好"
    # text = "应，进一步放大风险的危害，而传播环境的深刻变革对触发和放大生成式人工智 能传播风险也具有重要影响。生成式人工智能传播风险打破了独立、单层、定向的风险路径， 通过机内交互、人机交互与人人交互的复杂网络结构形成了联动、多层、多向的风险路径， 极易形成多领域共振传导。"
    result = get_chunk(database_name="XIAN",collection_name="ye85a36d295a199730340c5c755c1e2b",id="459963917155397121",text=text,resource="test.md")
    print(result)