#!/usr/bin/env python3
"""
简化版Memory Bank Web界面
"""

import streamlit as st
import asyncio
import sys
import os
from datetime import datetime

# 添加项目路径
sys.path.append('/Users/ben/liurenchaxin/src')

# 页面配置
st.set_page_config(
    page_title="Memory Bank 管理界面",
    page_icon="🧠",
    layout="wide"
)

# 标题
st.title("🧠 Vertex AI Memory Bank 管理界面")
st.markdown("---")

# 检查依赖
try:
    from jixia.memory.factory import get_memory_backend
    st.success("✅ 成功导入jixia模块")
except ImportError as e:
    st.error(f"❌ 无法导入jixia模块: {e}")
    st.info("请确保已激活虚拟环境并安装了所需依赖")
    st.stop()

# 八仙列表
EIGHT_IMMORTALS = [
    "lvdongbin", "tieguaili", "hanxiangzi", "lanzaihe",
    "hesengu", "zhonghanli", "caogujiu", "hanzhongli"
]

# 缓存Memory Bank后端
@st.cache_resource
def get_memory_backend_cached():
    """获取Memory Bank后端（缓存）"""
    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        backend = loop.run_until_complete(get_memory_backend("vertex"))
        return backend, loop
    except Exception as e:
        st.error(f"初始化Memory Bank失败: {e}")
        return None, None

# 异步函数包装器
def run_async(coro):
    """运行异步函数"""
    backend, loop = get_memory_backend_cached()
    if backend is None:
        return None
    try:
        return loop.run_until_complete(coro)
    except Exception as e:
        st.error(f"操作失败: {e}")
        return None

# 主界面选项卡
tab1, tab2, tab3 = st.tabs(["📋 Memory Bank列表", "🔍 搜索记忆", "➕ 添加记忆"])

with tab1:
    st.header("Memory Bank 实例列表")
    
    if st.button("🔄 刷新列表"):
        st.rerun()
    
    # 显示八仙Memory Bank状态
    cols = st.columns(4)
    for i, immortal in enumerate(EIGHT_IMMORTALS):
        with cols[i % 4]:
            st.subheader(f"🧙‍♂️ {immortal}")
            
            # 检查Memory Bank状态
            backend, _ = get_memory_backend_cached()
            if backend:
                try:
                    # 尝试获取agent context来验证Memory Bank存在
                    context = run_async(backend.get_agent_context(immortal))
                    if context is not None:
                        st.success("✅ 活跃")
                        
                        # 显示记忆数量
                        memories = run_async(backend.search_memories(immortal, "", limit=100))
                        if memories:
                            st.info(f"📝 记忆: {len(memories)}条")
                        else:
                            st.info("📝 记忆: 0条")
                    else:
                        st.warning("⚠️ 未初始化")
                except Exception as e:
                    st.error(f"❌ 错误")
            else:
                st.error("❌ 连接失败")

with tab2:
    st.header("🔍 搜索记忆")
    
    col1, col2 = st.columns([1, 2])
    
    with col1:
        selected_agent = st.selectbox("选择Agent", EIGHT_IMMORTALS)
        search_query = st.text_input("搜索关键词", placeholder="输入要搜索的内容...")
        
        if st.button("🔍 搜索"):
            if search_query:
                with st.spinner("搜索中..."):
                    backend, _ = get_memory_backend_cached()
                    if backend:
                        memories = run_async(backend.search_memories(selected_agent, search_query, limit=10))
                        st.session_state['search_results'] = memories
            else:
                st.warning("请输入搜索关键词")
    
    with col2:
        st.subheader("搜索结果")
        
        if 'search_results' in st.session_state and st.session_state['search_results']:
            st.success(f"找到 {len(st.session_state['search_results'])} 条记忆")
            
            for i, memory in enumerate(st.session_state['search_results']):
                with st.expander(f"记忆 {i+1}"):
                    st.write(f"**内容**: {memory.get('content', 'N/A')}")
                    st.write(f"**类型**: {memory.get('memory_type', 'N/A')}")
        elif 'search_results' in st.session_state:
            st.info("未找到匹配的记忆")
        else:
            st.info("请执行搜索以查看结果")

with tab3:
    st.header("➕ 添加记忆")
    
    add_agent = st.selectbox("选择Agent", EIGHT_IMMORTALS, key="add_agent")
    memory_type = st.selectbox("记忆类型", ["conversation", "preference", "knowledge", "strategy"])
    memory_content = st.text_area("记忆内容", placeholder="输入要添加的记忆内容...", height=150)
    
    if st.button("➕ 添加记忆"):
        if memory_content:
            with st.spinner("添加记忆中..."):
                backend, _ = get_memory_backend_cached()
                if backend:
                    # 准备元数据
                    metadata = {
                        "timestamp": datetime.now().isoformat(),
                        "source": "web_interface"
                    }
                    
                    # 添加记忆
                    success = run_async(backend.add_memory(
                        agent_id=add_agent,
                        content=memory_content,
                        memory_type=memory_type,
                        metadata=metadata
                    ))
                    
                    if success:
                        st.success("✅ 记忆添加成功！")
                    else:
                        st.error("❌ 添加记忆失败")
        else:
            st.warning("请输入记忆内容")

# 页脚
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: #666;'>
        🧠 Vertex AI Memory Bank Web界面
    </div>
    """, 
    unsafe_allow_html=True
)

# 使用说明
with st.expander("📖 使用说明"):
    st.markdown("""
    ### 功能说明
    
    1. **Memory Bank列表**: 查看所有八仙角色的Memory Bank状态和记忆数量
    2. **搜索记忆**: 在指定Agent的记忆中搜索特定内容
    3. **添加记忆**: 为Agent添加新的记忆
    
    ### 使用前准备
    
    1. 确保已激活虚拟环境: `source venv/bin/activate`
    2. 确保已设置Google Cloud认证: `gcloud auth application-default login`
    3. 运行此界面: `streamlit run simple_memory_web.py`
    """)