# -*- coding: utf-8 -*-
"""
Streamlit 前端 v2 - 增强版多智能体实时协作可视化界面
功能：
1. Agent 协作流程图（桑基图）
2. 实时状态监控面板
3. 分屏对话展示
4. 生成的文件预览
5. 时间线视图
"""
import streamlit as st
import os
import sys
from pathlib import Path
from datetime import datetime
import time
from typing import Dict, Any, List
import json

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

try:
    from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager
    AUTOGEN_AVAILABLE = True
except ImportError:
    AUTOGEN_AVAILABLE = False

from config.llm_config import (
    get_llm_config, 
    PM_PROMPT, 
    QA_PROMPT, 
    DEV_PROMPT, 
    ORCH_PROMPT
)


# 页面配置
st.set_page_config(
    page_title="AutoGen SDLC 多智能体协作平台",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义 CSS 样式
st.markdown("""
<style>
/* Agent 消息样式 */
.agent-message {
    padding: 1rem;
    border-radius: 0.5rem;
    margin-bottom: 1rem;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.pm-agent { background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%); border-left: 5px solid #2196f3; }
.qa-agent { background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%); border-left: 5px solid #4caf50; }
.dev-agent { background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%); border-left: 5px solid #ff9800; }
.orchestrator { background: linear-gradient(135deg, #f3e5f5 0%, #e1bee7 100%); border-left: 5px solid #9c27b0; }
.system { background: linear-gradient(135deg, #fce4ec 0%, #f8bbd9 100%); border-left: 5px solid #e91e63; }
.user { background: linear-gradient(135deg, #efebe9 0%, #d7ccc8 100%); border-left: 5px solid #795548; }

/* 状态卡片 */
.status-card {
    padding: 1.5rem;
    border-radius: 0.5rem;
    text-align: center;
    color: white;
    font-weight: bold;
}
.status-running { background: linear-gradient(135deg, #4caf50, #8bc34a); }
.status-stopped { background: linear-gradient(135deg, #f44336, #e91e63); }
.status-pending { background: linear-gradient(135deg, #ff9800, #ffc107); }

/* Agent 状态指示器 */
.agent-status-indicator {
    display: inline-block;
    width: 12px;
    height: 12px;
    border-radius: 50%;
    margin-right: 8px;
    animation: pulse 2s infinite;
}
@keyframes pulse {
    0% { opacity: 1; }
    50% { opacity: 0.5; }
    100% { opacity: 1; }
}

/* 文件卡片 */
.file-card {
    padding: 1rem;
    border: 1px solid #ddd;
    border-radius: 0.5rem;
    margin-bottom: 0.5rem;
    background: #f9f9f9;
}

/* 时间线 */
.timeline-item {
    padding: 1rem;
    border-left: 3px solid #2196f3;
    margin-left: 1rem;
    position: relative;
}
.timeline-item::before {
    content: '';
    position: absolute;
    left: -8px;
    top: 1rem;
    width: 13px;
    height: 13px;
    border-radius: 50%;
    background: #2196f3;
}
</style>
""", unsafe_allow_html=True)


# Agent 配置信息
AGENT_CONFIG = {
    "PM_Agent": {
        "name": "产品经理",
        "avatar": "📋",
        "color": "#2196f3",
        "description": "需求分析与 SRS 生成",
        "icon": "📊"
    },
    "QA_Agent": {
        "name": "测试工程师",
        "avatar": "✅",
        "color": "#4caf50",
        "description": "测试用例设计与 TDD",
        "icon": "🧪"
    },
    "Dev_Agent": {
        "name": "开发工程师",
        "avatar": "💻",
        "color": "#ff9800",
        "description": "代码实现与修复",
        "icon": "⚙️"
    },
    "Orchestrator": {
        "name": "协调器",
        "avatar": "🎯",
        "color": "#9c27b0",
        "description": "流程调度与验证",
        "icon": "🎭"
    },
    "User_Proxy": {
        "name": "用户代理",
        "avatar": "👤",
        "color": "#795548",
        "description": "人机交互接口",
        "icon": "🔌"
    }
}


def init_session_state():
    """初始化 session state"""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "is_running" not in st.session_state:
        st.session_state.is_running = False
    if "workflow_result" not in st.session_state:
        st.session_state.workflow_result = None
    if "conversation_history" not in st.session_state:
        st.session_state.conversation_history = []
    if "current_step" not in st.session_state:
        st.session_state.current_step = 0
    if "agent_stats" not in st.session_state:
        st.session_state.agent_stats = {
            "PM_Agent": {"messages": 0, "last_active": None},
            "QA_Agent": {"messages": 0, "last_active": None},
            "Dev_Agent": {"messages": 0, "last_active": None},
            "Orchestrator": {"messages": 0, "last_active": None},
            "User_Proxy": {"messages": 0, "last_active": None}
        }
    if "generated_files" not in st.session_state:
        st.session_state.generated_files = []
    if "active_agent" not in st.session_state:
        st.session_state.active_agent = None


def create_agents(api_key: str, base_url: str, model: str):
    """创建所有 Agent"""
    llm_config = get_llm_config(model=model, api_key=api_key, base_url=base_url)
    
    pm_agent = AssistantAgent(
        name="PM_Agent",
        system_message=PM_PROMPT,
        llm_config=llm_config,
        description="资深软件产品经理",
        human_input_mode="NEVER"
    )
    
    qa_agent = AssistantAgent(
        name="QA_Agent",
        system_message=QA_PROMPT,
        llm_config=llm_config,
        description="资深测试工程师",
        human_input_mode="NEVER"
    )
    
    dev_agent = AssistantAgent(
        name="Dev_Agent",
        system_message=DEV_PROMPT,
        llm_config=llm_config,
        description="资深软件工程师",
        human_input_mode="NEVER"
    )
    
    orchestrator = AssistantAgent(
        name="Orchestrator",
        system_message=ORCH_PROMPT,
        llm_config=llm_config,
        description="多智能体协调器",
        human_input_mode="NEVER"
    )
    
    user_proxy = UserProxyAgent(
        name="User_Proxy",
        human_input_mode="NEVER",  # 修复：Web 环境不支持 TERMINAL
        max_consecutive_auto_reply=0,
        code_execution_config={
            "work_dir": "workspace",
            "use_docker": False,
        },
        is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE")
    )
    
    return pm_agent, qa_agent, dev_agent, orchestrator, user_proxy, llm_config


def display_agent_status_panel():
    """显示 Agent 状态面板"""
    st.subheader("🎭 Agent 实时状态")
    
    cols = st.columns(5)
    agent_names = list(AGENT_CONFIG.keys())
    
    for i, (agent_name, config) in enumerate(AGENT_CONFIG.items()):
        with cols[i]:
            stats = st.session_state.agent_stats.get(agent_name, {})
            msg_count = stats.get("messages", 0)
            last_active = stats.get("last_active")
            
            is_active = st.session_state.active_agent == agent_name
            
            status_color = config["color"]
            if is_active:
                status_html = f'<span class="agent-status-indicator" style="background: {status_color}; animation: pulse 1s infinite;"></span>'
            else:
                status_html = f'<span style="display:inline-block;width:12px;height:12px;border-radius:50%;background:{status_color};opacity:0.3;"></span>'
            
            st.markdown(f"""
            <div style="text-align:center;padding:1rem;border:2px {'solid' if is_active else 'dashed'} {status_color};border-radius:0.5rem;">
                <div style="font-size:2rem;">{config['avatar']}</div>
                <div style="font-weight:bold;color:{status_color};">{config['name']}</div>
                <div style="font-size:0.8rem;color:#666;">消息：{msg_count}</div>
                <div style="font-size:0.7rem;color:#999;">{status_html}{'工作中' if is_active else '等待中'}</div>
            </div>
            """, unsafe_allow_html=True)


def display_workflow_diagram():
    """显示工作流程图"""
    st.subheader("🔄 工作流程图")
    
    # 使用 Mermaid 绘制流程图
    mermaid_code = """
    ```mermaid
    graph LR
        A[👤 用户需求] --> B[📋 PM Agent]
        B --> C[📄 SRS 文档]
        C --> D[✅ QA Agent]
        D --> E[🧪 测试用例]
        E --> F[💻 Dev Agent]
        F --> G[💾 源代码]
        G --> H[🔧 自动测试]
        H -->|失败 | F
        H -->|通过 | I[🎯 Orchestrator]
        I --> J[✅ 最终报告]
        
        style B fill:#e3f2fd,stroke:#2196f3,stroke-width:3px
        style D fill:#e8f5e9,stroke:#4caf50,stroke-width:3px
        style F fill:#fff3e0,stroke:#ff9800,stroke-width:3px
        style I fill:#f3e5f5,stroke:#9c27b0,stroke-width:3px
    ```
    """
    st.markdown(mermaid_code)


def display_message(agent_name: str, message: str, timestamp: str, index: int):
    """显示单条消息（增强版）"""
    config = AGENT_CONFIG.get(agent_name, {"name": agent_name, "avatar": "🤖", "color": "#999"})
    
    # 折叠长消息
    if len(message) > 500:
        with st.expander(f"{config['avatar']} {config['name']} - {timestamp[:16]}"):
            st.markdown(message)
    else:
        with st.chat_message(name=agent_name.lower().replace("_", ""), avatar=config['avatar']):
            st.markdown(f"**{config['avatar']} {config['name']}** *({timestamp[:16]})*")
            st.markdown(message)


def display_timeline():
    """显示时间线视图"""
    st.subheader("⏱️ 执行时间线")
    
    if not st.session_state.messages:
        st.info("暂无执行记录")
        return
    
    # 按时间分组显示关键事件
    timeline_data = []
    for msg in st.session_state.messages:
        agent = msg.get("agent_name", "Unknown")
        timestamp = msg.get("timestamp", "")[:19].replace('T', ' ')
        message = msg.get("message", "")[:100]
        
        if agent in AGENT_CONFIG:
            timeline_data.append({
                "time": timestamp,
                "agent": agent,
                "message": message,
                "icon": AGENT_CONFIG[agent]["avatar"]
            })
    
    # 倒序显示（最新的在前）
    for item in reversed(timeline_data[-10:]):  # 只显示最近 10 条
        with st.container():
            st.markdown(f"""
            <div class="timeline-item">
                <strong>{item['icon']} {item['time']}</strong><br>
                <span style="color:#666;">{AGENT_CONFIG.get(item['agent'], {}).get('name', item['agent'])}</span><br>
                {item['message']}...
            </div>
            """, unsafe_allow_html=True)


def display_generated_files():
    """显示生成的文件"""
    st.subheader("📁 生成的文件")
    
    workspace_dir = Path("workspace")
    if not workspace_dir.exists():
        st.info("工作目录为空")
        return
    
    files = list(workspace_dir.glob("*"))
    
    if not files:
        st.info("暂无生成的文件")
        return
    
    for file_path in files:
        if file_path.is_file():
            file_size = file_path.stat().st_size
            file_type = file_path.suffix
            
            icon_map = {
                ".py": "🐍",
                ".md": "📝",
                ".txt": "📄",
                ".json": "📊"
            }
            
            icon = icon_map.get(file_type, "📎")
            
            with st.container():
                st.markdown(f"""
                <div class="file-card">
                    <div style="display:flex;justify-content:space-between;align-items:center;">
                        <div>
                            <span style="font-size:1.5rem;">{icon}</span>
                            <strong>{file_path.name}</strong>
                            <span style="color:#666;font-size:0.9rem;margin-left:1rem;">
                                {file_size:,} bytes
                            </span>
                        </div>
                    </div>
                </div>
                """, unsafe_allow_html=True)
                
                # 文件内容预览
                if st.button(f"👁️ 预览 {file_path.name}", key=f"preview_{file_path.name}"):
                    try:
                        with open(file_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                        st.code(content[:2000] + ("..." if len(content) > 2000 else ""))
                    except Exception as e:
                        st.error(f"读取失败：{e}")


def display_chat_interface():
    """显示聊天界面"""
    st.subheader("💬 Agent 对话实录")
    
    if not st.session_state.messages:
        st.info("👈 暂无对话，请在下方输入需求并启动工作流")
    else:
        # 使用容器显示消息
        chat_container = st.container()
        with chat_container:
            for idx, msg in enumerate(st.session_state.messages):
                display_message(
                    agent_name=msg.get("agent_name", "Unknown"),
                    message=msg.get("message", ""),
                    timestamp=msg.get("timestamp", ""),
                    index=idx
                )


def export_conversation(messages: List[Dict]) -> str:
    """导出对话为 Markdown"""
    md_content = "# AutoGen SDLC 对话历史\n\n"
    md_content += f"**导出时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    md_content += "---\n\n"
    
    for msg in messages:
        timestamp = msg.get("timestamp", "")[:19].replace('T', ' ')
        agent = msg.get("agent_name", "Unknown")
        content = msg.get("message", "")
        
        config = AGENT_CONFIG.get(agent, {"avatar": "🤖", "name": agent})
        
        md_content += f"### {config['avatar']} [{timestamp}] {config['name']}\n\n"
        md_content += f"{content}\n\n"
        md_content += "---\n\n"
    
    return md_content


def main():
    """主应用"""
    init_session_state()
    
    # ===== 顶部标题区 =====
    st.title("🤖 AutoGen SDLC 多智能体协作平台")
    st.markdown("**基于 AutoGen + Qwen3.5-flash 的端到端软件交付系统 · 实时可视化 Agent 协作**")
    
    # ===== 侧边栏 - 配置区域 =====
    with st.sidebar:
        st.title("⚙️ 控制中心")
        
        # API 配置
        st.subheader("🔑 模型配置")
        api_key = st.text_input(
            "API Key",
            type="password",
            value=os.getenv("DASHSCOPE_API_KEY", ""),
            help="阿里云 DashScope API Key"
        )
        
        base_url = st.text_input(
            "Base URL",
            value="https://dashscope.aliyuncs.com/compatible-mode/v1",
            help="API Base URL"
        )
        
        model = st.selectbox(
            "模型选择",
            options=["qwen3.5-flash", "qwen-max", "qwen-plus", "qwen-turbo"],
            index=0,
            help="选择使用的 Qwen 模型"
        )
        
        # Agent 参数
        st.subheader("🎛️ Agent 参数")
        max_round = st.slider(
            "最大对话轮数",
            min_value=5,
            max_value=50,
            value=20,
            step=1
        )
        
        temperature = st.slider(
            "温度参数",
            min_value=0.0,
            max_value=1.0,
            value=0.7,
            step=0.1
        )
        
        st.divider()
        
        # 控制按钮
        st.subheader("🎮 流程控制")
        col1, col2 = st.columns(2)
        
        with col1:
            start_btn = st.button("▶️ 启动", use_container_width=True, type="primary")
        with col2:
            stop_btn = st.button("⏸️ 暂停", use_container_width=True)
        
        st.divider()
        
        # 导出选项
        st.subheader("💾 数据导出")
        if st.button("📄 导出 JSON", use_container_width=True):
            if st.session_state.messages:
                json_str = json.dumps(st.session_state.messages, ensure_ascii=False, indent=2)
                st.download_button(
                    label="下载 JSON",
                    data=json_str,
                    file_name=f"conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                    mime="application/json"
                )
        
        if st.button("📝 导出 Markdown", use_container_width=True):
            if st.session_state.messages:
                md_content = export_conversation(st.session_state.messages)
                st.download_button(
                    label="下载 Markdown",
                    data=md_content,
                    file_name=f"conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md",
                    mime="text/markdown"
                )
    
    # ===== 主区域 - 三栏布局 =====
    # 第一行：状态指标
    status_col1, status_col2, status_col3, status_col4 = st.columns(4)
    
    with status_col1:
        total_msgs = len(st.session_state.messages)
        st.metric("💬 总消息数", total_msgs)
    
    with status_col2:
        status = "🟢 运行中" if st.session_state.is_running else "🔴 已停止"
        status_class = "running" if st.session_state.is_running else "stopped"
        st.markdown(f"""
        <div class="status-card status-{status_class}">
            {status}
        </div>
        """, unsafe_allow_html=True)
    
    with status_col3:
        result_status = "✅ 成功" if st.session_state.workflow_result and st.session_state.workflow_result.get("success") else "⏳ 未开始"
        st.metric("📊 工作流状态", result_status)
    
    with status_col4:
        files_count = len(list(Path("workspace").glob("*"))) if Path("workspace").exists() else 0
        st.metric("📁 生成文件数", files_count)
    
    st.divider()
    
    # 第二行：左侧 - Agent 状态 + 流程图，右侧 - 时间线
    left_col, right_col = st.columns([2, 1])
    
    with left_col:
        display_agent_status_panel()
        st.divider()
        display_workflow_diagram()
    
    with right_col:
        display_timeline()
    
    st.divider()
    
    # 第三行：聊天界面
    display_chat_interface()
    
    # ===== 底部 - 用户需求输入 =====
    st.divider()
    
    # 用户输入
    user_input = st.chat_input("💡 请输入您的软件需求，例如：我需要一个电池健康状态预测 API...")
    
    if user_input:
        # 添加用户消息
        user_msg = {
            "timestamp": datetime.now().isoformat(),
            "agent_name": "User_Proxy",
            "role": "user",
            "message": user_input
        }
        st.session_state.messages.append(user_msg)
        st.session_state.conversation_history.append(user_msg)
        
        # 更新统计
        st.session_state.agent_stats["User_Proxy"]["messages"] += 1
        st.session_state.agent_stats["User_Proxy"]["last_active"] = datetime.now().isoformat()
        
        # 刷新显示
        st.rerun()
    
    # ===== 工作流执行逻辑 =====
    if start_btn:
        if not api_key:
            st.error("请先设置 API Key！")
            st.stop()
        
        if not AUTOGEN_AVAILABLE:
            st.error("请先安装 AutoGen: pip install pyautogen")
            st.stop()
        
        # 获取最后一条用户消息
        user_messages = [m for m in st.session_state.messages if m.get("role") == "user"]
        if not user_messages:
            st.warning("请先输入需求！")
            st.stop()
        
        latest_requirement = user_messages[-1]["message"]
        
        # 设置运行状态
        st.session_state.is_running = True
        st.session_state.current_step = 0
        
        # 创建进度占位符
        progress_placeholder = st.empty()
        
        with progress_placeholder:
            st.info("🚀 正在启动 SDLC 工作流，请稍候...")
        
        try:
            # 创建 Agent
            pm_agent, qa_agent, dev_agent, orchestrator, user_proxy, llm_config = create_agents(
                api_key=api_key,
                base_url=base_url,
                model=model
            )
            
            # 创建 GroupChat
            groupchat = GroupChat(
                agents=[pm_agent, qa_agent, dev_agent, orchestrator, user_proxy],
                messages=[],
                max_round=max_round,
                speaker_selection_method="round_robin"
            )
            
            manager = GroupChatManager(groupchat=groupchat, llm_config=llm_config)
            
            # 构建初始消息
            initial_message = f"""
请启动完整的 SDLC 流程：

【用户需求】
{latest_requirement}

【工作流程】
1. PM_Agent → 生成 SRS
2. QA_Agent → 生成测试用例
3. Dev_Agent → 编写代码
4. User_Proxy → 执行测试
5. Orchestrator → 汇总报告

开始协作！
"""
            
            # 记录第一条消息
            first_msg = {
                "timestamp": datetime.now().isoformat(),
                "agent_name": "Orchestrator",
                "role": "assistant",
                "message": f"🎯 启动 SDLC 工作流，需求：{latest_requirement[:100]}..."
            }
            st.session_state.messages.append(first_msg)
            
            # 更新统计
            st.session_state.agent_stats["Orchestrator"]["messages"] += 1
            st.session_state.active_agent = "Orchestrator"
            
            # 执行对话
            chat_result = user_proxy.initiate_chat(
                manager,
                message=initial_message,
                max_turns=max_round
            )
            
            # 记录所有对话
            for msg in groupchat.messages:
                chat_msg = {
                    "timestamp": datetime.now().isoformat(),
                    "agent_name": msg.get("name", "Unknown"),
                    "role": msg.get("role", "assistant"),
                    "message": msg.get("content", "")
                }
                
                # 避免重复添加
                if chat_msg not in st.session_state.messages:
                    st.session_state.messages.append(chat_msg)
                    
                    # 更新统计
                    agent_name = msg.get("name", "Unknown")
                    if agent_name in st.session_state.agent_stats:
                        st.session_state.agent_stats[agent_name]["messages"] += 1
                        st.session_state.agent_stats[agent_name]["last_active"] = datetime.now().isoformat()
            
            # 扫描生成的文件
            if Path("workspace").exists():
                files = list(Path("workspace").glob("*"))
                st.session_state.generated_files = [str(f) for f in files if f.is_file()]
            
            # 更新状态
            st.session_state.workflow_result = {
                "success": True,
                "summary": chat_result.summary if hasattr(chat_result, 'summary') else "完成"
            }
            st.session_state.current_step = 5
            st.session_state.is_running = False
            st.session_state.active_agent = None
            
            progress_placeholder.success("✅ SDLC 工作流完成！")
            
            # 刷新显示
            st.rerun()
            
        except Exception as e:
            st.session_state.is_running = False
            st.session_state.active_agent = None
            st.error(f"❌ 错误：{str(e)}")
            st.session_state.workflow_result = {"success": False, "error": str(e)}
    
    if stop_btn:
        st.session_state.is_running = False
        st.session_state.active_agent = None
        st.info("⏸️ 工作流已暂停")


if __name__ == "__main__":
    main()