#!/usr/bin/env python3
"""
Vertex AI Memory Bank Web界面
一个简单的Streamlit应用，用于通过Web界面访问和管理Memory Bank
"""

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

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

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

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

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

# 侧边栏配置
st.sidebar.header("配置")
project_id = st.sidebar.text_input("项目ID", value="inner-radius-469712-e9")
location = st.sidebar.text_input("区域", value="us-central1")

# 八仙列表
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, tab4 = st.tabs(["📋 Memory Bank列表", "🔍 搜索记忆", "➕ 添加记忆", "📊 统计信息"])

with tab1:
    st.header("Memory Bank 实例列表")
    
    if st.button("🔄 刷新列表", key="refresh_list"):
        st.rerun()
    
    # 显示八仙Memory Bank状态
    cols = st.columns(4)
    for i, immortal in enumerate(EIGHT_IMMORTALS):
        with cols[i % 4]:
            with st.container():
                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"❌ 错误: {str(e)[:50]}...")
                else:
                    st.error("❌ 连接失败")

with tab2:
    st.header("🔍 搜索记忆")
    
    col1, col2 = st.columns([1, 2])
    
    with col1:
        selected_agent = st.selectbox("选择Agent", EIGHT_IMMORTALS, key="search_agent")
        search_query = st.text_input("搜索关键词", placeholder="输入要搜索的内容...", key="search_query")
        search_limit = st.slider("结果数量", 1, 50, 10, key="search_limit")
        
        if st.button("🔍 搜索", key="search_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=search_limit))
                        st.session_state['search_results'] = memories
                        st.session_state['search_agent'] = selected_agent
                        st.session_state['search_query'] = search_query
            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}: {memory.get('content', 'N/A')[:50]}..."):
                    st.write(f"**内容**: {memory.get('content', 'N/A')}")
                    st.write(f"**类型**: {memory.get('memory_type', 'N/A')}")
                    st.write(f"**时间**: {memory.get('timestamp', 'N/A')}")
                    if 'metadata' in memory:
                        st.write(f"**元数据**: {memory['metadata']}")
        elif 'search_results' in st.session_state:
            st.info("未找到匹配的记忆")
        else:
            st.info("请执行搜索以查看结果")

with tab3:
    st.header("➕ 添加记忆")
    
    col1, col2 = st.columns([1, 1])
    
    with col1:
        add_agent = st.selectbox("选择Agent", EIGHT_IMMORTALS, key="add_agent")
        memory_type = st.selectbox("记忆类型", ["conversation", "preference", "knowledge", "strategy"], key="memory_type")
        memory_content = st.text_area("记忆内容", placeholder="输入要添加的记忆内容...", height=150, key="memory_content")
        
        # 可选的元数据
        st.subheader("元数据（可选）")
        importance = st.slider("重要性", 1, 10, 5, key="importance")
        tags = st.text_input("标签（用逗号分隔）", placeholder="标签1, 标签2, 标签3", key="tags")
        
        if st.button("➕ 添加记忆", key="add_memory_button"):
            if memory_content:
                with st.spinner("添加记忆中..."):
                    backend, _ = get_memory_backend_cached()
                    if backend:
                        # 准备元数据
                        metadata = {
                            "importance": importance,
                            "timestamp": datetime.now().isoformat(),
                            "source": "web_interface"
                        }
                        if tags:
                            metadata["tags"] = [tag.strip() for tag in tags.split(",")]
                        
                        # 添加记忆
                        success = run_async(backend.add_memory(
                            agent_id=add_agent,
                            content=memory_content,
                            memory_type=memory_type,
                            metadata=metadata
                        ))
                        
                        if success:
                            st.success("✅ 记忆添加成功！")
                            # 清空输入
                            st.session_state['memory_content'] = ""
                            st.session_state['tags'] = ""
                        else:
                            st.error("❌ 添加记忆失败")
            else:
                st.warning("请输入记忆内容")
    
    with col2:
        st.subheader("添加记忆预览")
        if memory_content:
            st.info(f"**Agent**: {add_agent}")
            st.info(f"**类型**: {memory_type}")
            st.info(f"**内容**: {memory_content}")
            st.info(f"**重要性**: {importance}/10")
            if tags:
                st.info(f"**标签**: {tags}")
        else:
            st.info("输入记忆内容以查看预览")

with tab4:
    st.header("📊 统计信息")
    
    if st.button("🔄 刷新统计", key="refresh_stats"):
        st.rerun()
    
    # 获取统计信息
    backend, _ = get_memory_backend_cached()
    if backend:
        stats_data = []
        
        for immortal in EIGHT_IMMORTALS:
            try:
                # 获取记忆数量
                memories = run_async(backend.search_memories(immortal, "", limit=1000))
                memory_count = len(memories) if memories else 0
                
                # 获取agent context
                context = run_async(backend.get_agent_context(immortal))
                status = "活跃" if context else "未初始化"
                
                stats_data.append({
                    "Agent": immortal,
                    "状态": status,
                    "记忆数量": memory_count,
                    "最后更新": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                })
            except Exception as e:
                stats_data.append({
                    "Agent": immortal,
                    "状态": "错误",
                    "记忆数量": 0,
                    "最后更新": f"错误: {str(e)[:30]}..."
                })
        
        # 显示统计表格
        st.dataframe(stats_data, use_container_width=True)
        
        # 显示汇总信息
        col1, col2, col3, col4 = st.columns(4)
        
        total_agents = len(EIGHT_IMMORTALS)
        active_agents = sum(1 for item in stats_data if item["状态"] == "活跃")
        total_memories = sum(item["记忆数量"] for item in stats_data)
        avg_memories = total_memories / total_agents if total_agents > 0 else 0
        
        with col1:
            st.metric("总Agent数", total_agents)
        
        with col2:
            st.metric("活跃Agent数", active_agents)
        
        with col3:
            st.metric("总记忆数", total_memories)
        
        with col4:
            st.metric("平均记忆数", f"{avg_memories:.1f}")

# 页脚
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: #666;'>
        🧠 Vertex AI Memory Bank Web界面 | 
        <a href='https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/memory-bank/overview' target='_blank'>官方文档</a>
    </div>
    """, 
    unsafe_allow_html=True
)

# 使用说明
with st.expander("📖 使用说明"):
    st.markdown("""
    ### 功能说明
    
    1. **Memory Bank列表**: 查看所有八仙角色的Memory Bank状态和记忆数量
    2. **搜索记忆**: 在指定Agent的记忆中搜索特定内容
    3. **添加记忆**: 为Agent添加新的记忆，支持不同类型和元数据
    4. **统计信息**: 查看所有Agent的统计数据和汇总信息
    
    ### 使用前准备
    
    1. 确保已激活虚拟环境: `source venv/bin/activate`
    2. 确保已设置Google Cloud认证: `gcloud auth application-default login`
    3. 运行此界面: `streamlit run memory_bank_web_interface.py`
    
    ### 注意事项
    
    - Memory Bank目前仅在us-central1区域可用
    - 搜索功能支持模糊匹配
    - 添加的记忆会立即生效
    - 统计信息实时更新
    """)