#!/usr/bin/env python3
"""
启用bbav5持久化记忆
"""
import asyncio
from pathlib import Path

from gfa_core.mas.memory import MemoryManager
from gfa_core.mas.memory.persistent import PersistentMemory
from gfa_core.mas.memory.session import SessionMemory, SessionData
from gfa_core.mas.memory.case import CaseMemory, AnalysisCase
from gfa_core.mas.memory.knowledge import KnowledgeMemory, KnowledgeEntry


async def create_persistent_memory_manager(storage_dir: Path = None) -> MemoryManager:
    """创建带持久化的MemoryManager"""
    
    if storage_dir is None:
        storage_dir = Path("/home/xia-lab/wutx/00test/bbav5/memory_storage")
    
    storage_dir.mkdir(parents=True, exist_ok=True)
    
    # 创建持久化存储
    session_storage = PersistentMemory(storage_dir / "sessions")
    case_storage = PersistentMemory(storage_dir / "cases")
    knowledge_storage = PersistentMemory(storage_dir / "knowledge")
    
    # 创建MemoryManager
    manager = MemoryManager(
        session_memory=SessionMemory(),  # 内存缓存 + 持久化备份
        case_memory=CaseMemory(),
        knowledge_memory=KnowledgeMemory(),
    )
    
    # 可选：添加持久化包装器
    # 或修改SessionMemory等使用PersistentMemory作为后端
    
    print(f"✅ MemoryManager创建成功")
    print(f"   存储目录: {storage_dir}")
    print(f"   Session: {storage_dir / 'sessions'}")
    print(f"   Case: {storage_dir / 'cases'}")
    print(f"   Knowledge: {storage_dir / 'knowledge'}")
    
    return manager


async def test_persistent_memory():
    """测试持久化记忆"""
    manager = await create_persistent_memory_manager()
    
    # 测试Session
    session_id = "test_session_001"
    session = await manager.create_session(session_id, {"test": "data"})
    print(f"\n✅ Session创建: {session_id}")
    
    # 添加消息
    await manager.add_message_to_session(session_id, "user", "Hello")
    await manager.add_message_to_session(session_id, "assistant", "Hi there")
    print(f"✅ 消息添加完成")
    
    # 测试Case
    case = await manager.create_case(
        case_id="test_case_001",
        user_id="user_001",
        gene_family="SPL",
        species="Arabidopsis",
    )
    print(f"✅ Case创建: {case.case_id}")
    
    # 更新结果
    await manager.update_case_results("test_case_001", {"status": "completed"})
    print(f"✅ Case结果更新")
    
    # 检索验证
    retrieved_session = await manager.get_session(session_id)
    retrieved_case = await manager.get_case("test_case_001")
    
    print(f"\n✅ 检索验证:")
    print(f"   Session: {retrieved_session.session_id if retrieved_session else 'None'}")
    print(f"   Case: {retrieved_case.case_id if retrieved_case else 'None'}")
    
    print("\n🎉 持久化记忆测试完成！")


if __name__ == "__main__":
    asyncio.run(test_persistent_memory())