app.py

import streamlit as st
import os
import uuid
import shutil
from PIL import Image
from main import run_analysis_logic

# --- 1. 页面基本配置 ---
st.set_page_config(
    page_title="DataAgent Pro (Multi-File Edition)",
    page_icon="🤖",
    layout="wide"
)

# --- 2. 状态初始化与临时目录管理 ---
if 'session_id' not in st.session_state:
    st.session_state['session_id'] = str(uuid.uuid4())

# 定义本次会话的临时数据目录 (绝对路径)
BASE_TEMP_DIR = os.path.join(os.getcwd(), "temp_storage")
SESSION_DATA_DIR = os.path.join(BASE_TEMP_DIR, st.session_state['session_id'])

# 确保目录存在
if not os.path.exists(SESSION_DATA_DIR):
    os.makedirs(SESSION_DATA_DIR, exist_ok=True)

# 初始化聊天记录
if "messages" not in st.session_state:
    st.session_state.messages = []

# --- 3. 侧边栏：文件管理 ---
with st.sidebar:
    st.title("📂 数据管理中心")
    st.markdown("---")

    # 上传组件：开启多文件上传
    uploaded_files = st.file_uploader(
        "上传 CSV 或 Excel (支持多选)",
        type=["csv", "xlsx"],
        accept_multiple_files=True
    )

    if uploaded_files:
        # 清空旧文件以便重新同步（可选，根据需求决定是否保留）
        # shutil.rmtree(SESSION_DATA_DIR)
        # os.makedirs(SESSION_DATA_DIR, exist_ok=True)

        for uploaded_file in uploaded_files:
            target_path = os.path.join(SESSION_DATA_DIR, uploaded_file.name)
            with open(target_path, "wb") as f:
                f.write(uploaded_file.getbuffer())

        st.success(f"✅ 已就绪 {len(uploaded_files)} 个文件")

        # 展示已识别的文件列表
        with st.expander("已加载文件列表"):
            for f in os.listdir(SESSION_DATA_DIR):
                st.text(f"📄 {f}")

    st.markdown("---")
    st.info("""
    **核心架构说明：**
    - **感知**: 自动探测目录下所有表的 Schema。
    - **决策**: LangGraph 编排多表关联逻辑。
    - **执行**: Docker 容器化物理隔离。
    """)

    if st.button("🗑️ 清空所有数据"):
        shutil.rmtree(SESSION_DATA_DIR)
        os.makedirs(SESSION_DATA_DIR, exist_ok=True)
        st.session_state.messages = []
        st.rerun()

# --- 4. 主界面：对话区 ---
st.title("🤖 DataAgent Pro: 多表关联分析")
st.caption("基于 LangGraph + Docker 的自进化数据智能体 | 支持跨表 Join 与 自动修复")

# 显示历史消息
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if "chart" in message:
            st.image(message["chart"])

# 用户输入
if query := st.chat_input("例如：关联房价表和城市信息表，分析不同等级城市的平均单价"):

    if not os.listdir(SESSION_DATA_DIR):
        st.warning("⚠️ 请先在左侧上传至少一个数据文件。")
        st.stop()

    # 1. 记录用户消息
    st.session_state.messages.append({"role": "user", "content": query})
    with st.chat_message("user"):
        st.markdown(query)

    # 2. Agent 执行
    with st.chat_message("assistant"):
        # 使用进度状态
        with st.status("Agent 正在处理多表任务...", expanded=True) as status:
            try:
                st.write("🔍 正在扫描目录并提取多表元数据...")
                # 注意：此时传递的是目录路径 SESSION_DATA_DIR
                final_state = run_analysis_logic(query, SESSION_DATA_DIR)

                st.write("🧠 正在规划关联路径并生成代码...")
                if final_state.get("retry_count", 0) > 0:
                    st.warning(f"🔄 发生了 {final_state['retry_count']} 次自修复重试")

                st.write("🐳 正在 Docker 沙箱中执行跨表计算...")

                if final_state.get("error"):
                    status.update(label="❌ 执行失败", state="error")
                else:
                    status.update(label="✅ 分析成功", state="complete")

            except Exception as e:
                st.error(f"系统运行异常: {str(e)}")
                status.update(label="💥 崩溃", state="error")
                st.stop()

        # 3. 展示结论
        res_logs = final_state.get("logs", "Agent 没有返回任何文字结论。")
        st.markdown("### 📊 分析结论")
        st.code(res_logs)

        # 4. 检查并展示图表
        # 注意：多文件模式下，图表通常保存在数据目录下
        chart_path = os.path.join(SESSION_DATA_DIR, "output_chart.png")
        chart_image = None

        if os.path.exists(chart_path):
            chart_image = Image.open(chart_path)
            st.image(chart_image, caption="多表关联分析图表")
            # 及时移除旧图，防止下次查询干扰（可选）
            # os.remove(chart_path)

        # 5. 保存历史记录
        ans_msg = {"role": "assistant", "content": res_logs}
        if chart_image:
            ans_msg["chart"] = chart_image
        st.session_state.messages.append(ans_msg)

# --- 5. 样式美化 ---
st.markdown("""
<style>
    .stCodeBlock {
        border-left: 5px solid #00c0f2;
    }
    .stChatMessage {
        border-radius: 10px;
    }
</style>
""", unsafe_allow_html=True)