import logging

from Volcengine.VikingDBMemoryService import MEMORY_COLLECTION_NAME, initialize_services, ensure_collection_exists
from Volcengine.chat import handle_conversation_turn, archive_conversation

# 控制日志输出
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 main():
    logger.info("开始端到端记忆测试...")

    try:
        memory_service, llm_client = initialize_services()
        collection_name = MEMORY_COLLECTION_NAME
        user_id = "system"
        assistant_id = "assistant"
        ensure_collection_exists(memory_service, collection_name)
    except Exception as e:
        logger.info(f"初始化失败: {e}")
        return

    logger.info("\n--- 阶段 1: 初始对话 ---")
    initial_conversation_history = []
    handle_conversation_turn(
        memory_service, llm_client, collection_name, user_id,
        "你好，我是小明，今年18岁，但压力好大。",
        initial_conversation_history
    )
    handle_conversation_turn(
        memory_service, llm_client, collection_name, user_id,
        "马上就要高考了，家里人的期待好高。",
        initial_conversation_history
    )

    logger.info("\n--- 阶段 2: 归档记忆 ---")
    archive_conversation(
        memory_service, collection_name, user_id, assistant_id,
        initial_conversation_history, "study_stress_discussion"
    )

    logger.info("\n--- 阶段 3: 验证记忆 ---")
    verification_conversation_history = []
    handle_conversation_turn(
        memory_service, llm_client, collection_name, user_id,
        "我最近很焦虑，不知道该怎么办。",
        verification_conversation_history
    )

    logger.info("\n端到端记忆测试完成！")


if __name__ == "__main__":
    main()