from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection,MilvusClient
from config import MILVUS_CLIENT
import time
from datetime import datetime, timedelta

def create_partition_by_hour(current_hour):
    # 连接到 Milvus 服务器
    connections.connect("default",uri=MILVUS_CLIENT)
    # 获取集合
    collection_name = "pdf_measure_v4"
    collection = Collection(collection_name)

    # 获取当前时间
    now = datetime.now()


    # 创建当前小时的分区
    partition_name = f"partition_{current_hour}"
    if not collection.has_partition(partition_name):
        collection.create_partition(partition_name)
        print(f"Created partition: {partition_name}")
        partition = collection.partition(partition_name)
        partition.load()

    # 删除前2个小时的分区
    previous_hour = (now - timedelta(hours=2)).strftime("%Y%m%d%H")
    previous_partition_name = f"partition_{previous_hour}"
    if collection.has_partition(previous_partition_name):

        pre_partition = collection.partition(previous_partition_name)
        pre_partition.release()
        collection.drop_partition(previous_partition_name)
        print(f"Dropped partition: {previous_partition_name}")





# data = []
# measure_data = {}
# vector = [0.61865162262130161] * 1536
# measure_data['vector'] = vector
# measure_data['table_num'] = int(2)
# measure_data['table_index'] = int(2)
# measure_data['measure_name'] = "234234"
# measure_data['measure_value'] = "23432"
# measure_data['measure_unit'] = "123423"
# measure_data['file_id'] = "100000"
#
# data.append(measure_data)
# res = client.insert(
#                 collection_name=collection_name,
#                 data=data,
#                 partition_name=partition_name
#             )

# filter_str = 'file_id == "'+"2122"+'"'
# res = client.search(
#     collection_name=collection_name,  # Replace with the actual name of your collection
#     # Replace with your query vector
#     data=data,
#     limit=3,  # Max. number of search results to return
#     search_params={"metric_type": "COSINE", "params": {}},  # Search parameters
#     output_fields=["measure_name", "measure_value", "table_num", "table_index", "measure_unit"],
#     filter=filter_str,
#     partition_name=partition_name
# )
# print(f"============================={res}")




# from pymilvus import connections, CollectionSchema, Collection,utility,FieldSchema,DataType
# # 连接到 B 服务器上的 Milvus
# # connections.connect(host='124.70.129.232', port='19530')# 测试服务器
# connections.connect(host='1.94.60.103', port='19530')# 测试服务器
# # # 获取集合列表
# utility.drop_collection("pdf_measure_v4")
#
# # 定义字段
# fields = [
#     FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
#     FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1536),
#     FieldSchema(name="table_num", dtype=DataType.INT16),
#     FieldSchema(name="table_index", dtype=DataType.INT16),
#     FieldSchema(name="measure_name", dtype=DataType.VARCHAR, max_length=200),
#     FieldSchema(name="measure_value", dtype=DataType.VARCHAR, max_length=200),
#     FieldSchema(name="file_id", dtype=DataType.VARCHAR, max_length=200),
#     FieldSchema(name="measure_unit", dtype=DataType.VARCHAR, max_length=200)
# ]
#
# # 定义集合的 schema
# schema = CollectionSchema(fields=fields, description="My Milvus collection")
#
# # 创建集合
# collection = Collection(name="pdf_measure_v4", schema=schema)
#
# collection = Collection("pdf_measure_v4")
# index_params = {
#     "index_type": "IVF_FLAT",
#     "metric_type": "COSINE",
#     "params": {"nlist": 128}
# }
# collection.create_index(field_name="vector", index_params=index_params)
# collection.load()