import json
import logging

from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME
from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY

# 控制日志输出
logger = logging.getLogger('CollectionMemory')
logger.setLevel(logging.INFO)

# 只添加一次处理器，避免重复日志
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(handler)

def drop_existing_collection(collection_name):
    # 初始化记忆库服务
    memory_service = VikingDBMemoryService(
        ak=VOLC_ACCESSKEY,
        sk=VOLC_SECRETKEY,
        host="api-knowledgebase.mlp.cn-beijing.volces.com",
        region="cn-beijing"
    )

    try:
        # 检查集合是否存在
        logger.info(f"正在检查集合 '{collection_name}'...")
        memory_service.get_collection(collection_name)
        logger.info(f"集合 '{collection_name}' 已存在，准备删除...")

        # 删除集合
        response = memory_service.drop_collection(collection_name)
        logger.info(f"删除响应: {json.dumps(response, ensure_ascii=False, indent=2)}")
        logger.info(f"集合 '{collection_name}' 删除成功")
        return True

    except Exception as e:
        error_msg = str(e)
        if "collection not exist" in error_msg:
            logger.info(f"集合 '{collection_name}' 不存在，无需删除")
            return False
        else:
            logger.error(f"操作失败: {error_msg}")
            raise

if __name__ == "__main__":
    # 可修改为目标集合名称
    drop_existing_collection(MEMORY_COLLECTION_NAME)