# -*- coding: utf-8 -*-
"""
Streamlit 前端 v3 - Agent 对话流可视化版本
核心特性：
1. 清晰展示每个 Agent 当前在做什么（任务标签）
2. Agent 之间的对话流（类似聊天室）
3. 实时活动指示器（高亮当前发言的 Agent）
4. 对话历史完整记录
"""
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 - Agent 对话流",
    page_icon="💬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义 CSS 样式 - 突出对话流
st.markdown("""
<style>
/* 主对话区域 */
.chat-flow-container {
    background: #f0f2f6;
    padding: 20px;
    border-radius: 10px;
    margin-bottom: 20px;
}

/* Agent 对话气泡 */
.agent-chat-bubble {
    padding: 15px 20px;
    margin: 10px 0;
    border-radius: 15px;
    max-width: 85%;
    position: relative;
    animation: bubble-in 0.3s ease;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

@keyframes bubble-in {
    from {
        opacity: 0;
        transform: translateY(10px);
    }
    to {
        opacity: 1;
        transform: translateY(0);
    }
}

/* 不同 Agent 的气泡样式 */
.pm-bubble { 
    background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%); 
    border: 2px solid #2196f3;
    margin-left: 0;
}

.qa-bubble { 
    background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%); 
    border: 2px solid #4caf50;
    margin-left: 0;
}

.dev-bubble { 
    background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%); 
    border: 2px solid #ff9800;
    margin-left: 0;
}

.orchestrator-bubble { 
    background: linear-gradient(135deg, #f3e5f5 0%, #e1bee7 100%); 
    border: 2px solid #9c27b0;
    margin-left: 0;
}

.user-bubble { 
    background: linear-gradient(135deg, #efebe9 0%, #d7ccc8 100%); 
    border: 2px solid #795548;
    margin-left: auto;
}

/* Agent 头像和名称 */
.agent-chat-header {
    display: flex;
    align-items: center;
    margin-bottom: 8px;
    font-weight: bold;
}

.agent-chat-avatar {
    font-size: 1.5rem;
    margin-right: 8px;
}

.agent-chat-name {
    font-size: 1rem;
    color: #333;
}

.agent-chat-time {
    margin-left: auto;
    font-size: 0.75rem;
    color: #666;
}

/* 任务标签 */
.task-badge {
    display: inline-block;
    padding: 3px 10px;
    border-radius: 12px;
    font-size: 0.75rem;
    font-weight: bold;
    margin-left: 8px;
    color: white;
}

.pm-task { background: #2196f3; }
.qa-task { background: #4caf50; }
.dev-task { background: #ff9800; }
.orchestrator-task { background: #9c27b0; }
.user-task { background: #795548; }

/* Agent 状态卡片 */
.agent-status-card {
    padding: 15px;
    border-radius: 10px;
    text-align: center;
    border: 3px solid #e0e0e0;
    transition: all 0.3s ease;
    background: white;
}

.agent-status-card.active {
    border-color: #4caf50;
    background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%);
    animation: agent-pulse 1.5s infinite;
}

@keyframes agent-pulse {
    0%, 100% { transform: scale(1); }
    50% { transform: scale(1.02); }
}

.agent-status-avatar {
    font-size: 2.5rem;
    margin-bottom: 5px;
}

.agent-status-name {
    font-weight: bold;
    font-size: 0.9rem;
    margin-bottom: 5px;
}

.agent-status-task {
    font-size: 0.75rem;
    color: #666;
    min-height: 20px;
}

.agent-status-indicator {
    display: inline-block;
    width: 8px;
    height: 8px;
    border-radius: 50%;
    margin-right: 5px;
}

.status-active {
    background: #4caf50;
    animation: dot-pulse 1s infinite;
}

.status-waiting {
    background: #ffc107;
}

.status-inactive {
    background: #e0e0e0;
}

@keyframes dot-pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.3; }
}

/* 对话流指示器 */
.chat-flow-indicator {
    display: flex;
    align-items: center;
    margin: 10px 0;
    padding: 10px;
    background: white;
    border-radius: 8px;
}

.flow-arrow {
    font-size: 1.5rem;
    color: #2196f3;
    margin: 0 10px;
}

/* 当前活动提示 */
.active-agent-banner {
    background: linear-gradient(135deg, #4caf50, #8bc34a);
    color: white;
    padding: 10px 20px;
    border-radius: 8px;
    text-align: center;
    font-weight: bold;
    margin-bottom: 20px;
    animation: banner-pulse 2s infinite;
}

@keyframes banner-pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.8; }
}
</style>
""", unsafe_allow_html=True)


# Agent 配置
AGENT_CONFIG = {
    "PM_Agent": {
        "name": "产品经理",
        "avatar": "📋",
        "color": "#2196f3",
        "bubble_class": "pm-bubble",
        "task_class": "pm-task",
        "default_task": "待命"
    },
    "QA_Agent": {
        "name": "测试工程师",
        "avatar": "✅",
        "color": "#4caf50",
        "bubble_class": "qa-bubble",
        "task_class": "qa-task",
        "default_task": "待命"
    },
    "Dev_Agent": {
        "name": "开发工程师",
        "avatar": "💻",
        "color": "#ff9800",
        "bubble_class": "dev-bubble",
        "task_class": "dev-task",
        "default_task": "待命"
    },
    "Orchestrator": {
        "name": "协调器",
        "avatar": "🎯",
        "color": "#9c27b0",
        "bubble_class": "orchestrator-bubble",
        "task_class": "orchestrator-task",
        "default_task": "待命"
    },
    "User_Proxy": {
        "name": "用户代理",
        "avatar": "👤",
        "color": "#795548",
        "bubble_class": "user-bubble",
        "task_class": "user-task",
        "default_task": "待命"
    }
}


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 "current_agent" not in st.session_state:
        st.session_state.current_agent = None
    if "agent_tasks" not in st.session_state:
        st.session_state.agent_tasks = {
            "PM_Agent": "待命",
            "QA_Agent": "待命",
            "Dev_Agent": "待命",
            "Orchestrator": "待命",
            "User_Proxy": "待命"
        }
    if "agent_stats" not in st.session_state:
        st.session_state.agent_stats = {agent: 0 for agent in AGENT_CONFIG}


def display_agent_status_row():
    """显示 Agent 状态行 - 突出当前活动的 Agent"""
    st.subheader("🎭 Agent 实时状态")
    
    cols = st.columns(5)
    
    for i, (agent_key, config) in enumerate(AGENT_CONFIG.items()):
        with cols[i]:
            is_active = st.session_state.current_agent == agent_key
            current_task = st.session_state.agent_tasks.get(agent_key, "待命")
            msg_count = st.session_state.agent_stats.get(agent_key, 0)
            
            status_class = "active" if is_active else ""
            indicator_class = "status-active" if is_active else "status-inactive"
            
            st.markdown(f"""
            <div class="agent-status-card {status_class}">
                <div class="agent-status-avatar">{config['avatar']}</div>
                <div class="agent-status-name">
                    <span class="agent-status-indicator {indicator_class}"></span>
                    {config['name']}
                </div>
                <div class="agent-status-task">
                    📝 {current_task}
                </div>
                <div style="font-size:0.7rem;color:#999;margin-top:5px;">
                    💬 {msg_count} 条消息
                </div>
            </div>
            """, unsafe_allow_html=True)


def display_active_agent_banner():
    """显示当前活动 Agent 的横幅"""
    if st.session_state.current_agent:
        config = AGENT_CONFIG.get(st.session_state.current_agent, {})
        st.markdown(f"""
        <div class="active-agent-banner">
            {config.get('avatar', '🤖')} 当前发言：{config.get('name', 'Unknown')} - {st.session_state.agent_tasks.get(st.session_state.current_agent, '工作中...')}
        </div>
        """, unsafe_allow_html=True)


def display_chat_flow():
    """显示 Agent 对话流 - 核心功能"""
    st.subheader("💬 Agent 对话流")
    
    if not st.session_state.messages:
        st.info("👈 暂无对话，请启动工作流")
        return
    
    # 使用容器显示对话流
    chat_container = st.container()
    
    for idx, msg in enumerate(st.session_state.messages):
        agent_key = msg.get("agent_key", "Unknown")
        config = AGENT_CONFIG.get(agent_key, {"name": "未知", "avatar": "🤖", "bubble_class": ""})
        task_class = AGENT_CONFIG.get(agent_key, {}).get("task_class", "")
        
        timestamp = msg.get("timestamp", "")
        if timestamp:
            time_str = datetime.fromisoformat(timestamp).strftime("%H:%M:%S")
        else:
            time_str = datetime.now().strftime("%H:%M:%S")
        
        content = msg.get("content", "")
        task = msg.get("task", "工作中")
        
        # 显示对话气泡
        st.markdown(f"""
        <div class="agent-chat-bubble {config['bubble_class']}">
            <div class="agent-chat-header">
                <span class="agent-chat-avatar">{config['avatar']}</span>
                <span class="agent-chat-name">{config['name']}</span>
                <span class="task-badge {task_class}">{task}</span>
                <span class="agent-chat-time">{time_str}</span>
            </div>
            <div style="color:#333;line-height:1.6;">{content}</div>
        </div>
        """, unsafe_allow_html=True)
        
        # 显示对话流指示器（下一条消息的箭头）
        if idx < len(st.session_state.messages) - 1:
            next_msg = st.session_state.messages[idx + 1]
            next_agent = next_msg.get("agent_key", "")
            next_config = AGENT_CONFIG.get(next_agent, {})
            
            if next_agent != agent_key:
                st.markdown(f"""
                <div class="chat-flow-indicator">
                    <div style="flex:1;"></div>
                    <div class="flow-arrow">⬇️</div>
                    <div style="flex:1;text-align:center;color:#666;font-size:0.8rem;">
                        传递给 {next_config.get('name', 'Unknown')}
                    </div>
                </div>
                """, unsafe_allow_html=True)


def save_generated_files():
    """从对话中提取并保存生成的文件到 workspace 目录"""
    workspace_dir = Path("workspace")
    workspace_dir.mkdir(parents=True, exist_ok=True)
    
    files_saved = []
    
    # 遍历所有消息，提取文件内容
    for msg in st.session_state.messages:
        content = msg.get("content", "")
        agent_key = msg.get("agent_key", "")
        
        # PM Agent 生成 SRS.md
        if agent_key == "PM_Agent" and ("SRS" in content or "软件需求规格" in content):
            srs_file = workspace_dir / "SRS.md"
            # 提取 SRS 内容（查找包含 FR- 或 NFR- 的段落）
            srs_content = content
            if "功能性需求" in content:
                srs_content = content[content.find("功能性需求"):]
            with open(srs_file, 'w', encoding='utf-8') as f:
                f.write(f"# 软件需求规格说明书 (SRS)\n\n")
                f.write(f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
                f.write(srs_content)
            files_saved.append(str(srs_file))
        
        # QA Agent 生成测试文件
        if agent_key == "QA_Agent" and ("test_" in content or "def test_" in content):
            test_file = workspace_dir / "test_battery_health.py"
            # 提取 Python 代码
            if "```python" in content:
                code = content.split("```python")[1].split("```")[0]
            else:
                code = content
            with open(test_file, 'w', encoding='utf-8') as f:
                f.write(f'"""\n电池健康状态测试用例\n生成时间：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n"""\n\n')
                f.write(code)
            files_saved.append(str(test_file))
        
        # Dev Agent 生成源代码
        if agent_key == "Dev_Agent" and ("def " in content or "class " in content):
            src_file = workspace_dir / "src_battery_health.py"
            # 提取 Python 代码
            if "```python" in content:
                code = content.split("```python")[1].split("```")[0]
            else:
                code = content
            with open(src_file, 'w', encoding='utf-8') as f:
                f.write(f'"""\n电池健康状态计算模块\n生成时间：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n"""\n\n')
                f.write(code)
            files_saved.append(str(src_file))
        
        # Orchestrator 生成最终报告
        if agent_key == "Orchestrator" and ("完成" in content and "SDLC" in content):
            report_file = workspace_dir / "FINAL_REPORT.md"
            with open(report_file, 'w', encoding='utf-8') as f:
                f.write(f"# SDLC 项目最终报告\n\n")
                f.write(f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
                f.write("## 项目概述\n\n")
                f.write("基于 AutoGen 多智能体系统完成的软件交付项目。\n\n")
                f.write("## 生成文件列表\n\n")
                for i, saved_file in enumerate(files_saved, 1):
                    f.write(f"{i}. {saved_file}\n")
                f.write(f"\n## 项目总结\n\n{content}\n")
            files_saved.append(str(report_file))
    
    return files_saved


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",
        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 add_message(agent_key: str, content: str, task: str = "工作中"):
    """添加消息到对话流"""
    msg = {
        "agent_key": agent_key,
        "content": content,
        "timestamp": datetime.now().isoformat(),
        "task": task
    }
    st.session_state.messages.append(msg)
    
    # 更新统计
    st.session_state.agent_stats[agent_key] = st.session_state.agent_stats.get(agent_key, 0) + 1
    
    # 更新当前 Agent
    st.session_state.current_agent = agent_key
    
    # 更新任务
    st.session_state.agent_tasks[agent_key] = task


def main():
    """主应用"""
    init_session_state()
    
    # 标题
    st.title("💬 AutoGen SDLC - Agent 对话流可视化")
    st.markdown("**清晰展示每个 Agent 在做什么 · 实时追踪 Agent 之间的对话交互**")
    
    # 侧边栏配置
    with st.sidebar:
        st.title("⚙️ 控制中心")
        
        api_key = st.text_input("API Key", type="password", value=os.getenv("DASHSCOPE_API_KEY", ""))
        base_url = st.text_input("Base URL", value="https://dashscope.aliyuncs.com/compatible-mode/v1")
        model = st.selectbox("模型选择", ["qwen3.5-flash", "qwen-max", "qwen-plus", "qwen-turbo"], index=0)
        
        max_round = st.slider("最大对话轮数", 5, 50, 20)
        
        st.divider()
        
        col1, col2 = st.columns(2)
        with col1:
            start_btn = st.button("▶️ 启动对话流", type="primary", use_container_width=True)
        with col2:
            stop_btn = st.button("⏸️ 暂停", use_container_width=True)
        
        st.divider()
        
        if st.button("🗑️ 清空对话", use_container_width=True):
            st.session_state.messages = []
            st.session_state.current_agent = None
            st.session_state.agent_tasks = {k: "待命" for k in AGENT_CONFIG}
            st.session_state.agent_stats = {k: 0 for k in AGENT_CONFIG}
            st.rerun()
        
        if st.button("📥 导出对话", 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"agent_chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                    mime="application/json"
                )
        
        st.divider()
        
        # 显示生成的文件
        st.subheader("📁 生成的文件")
        workspace_dir = Path("workspace")
        if workspace_dir.exists():
            files = list(workspace_dir.glob("*"))
            if files:
                for file in files:
                    if file.is_file():
                        with st.expander(f"📄 {file.name}"):
                            try:
                                content = file.read_text(encoding='utf-8')
                                st.code(content[:500] + ("..." if len(content) > 500 else ""))
                                st.download_button(
                                    label="⬇️ 下载",
                                    data=content,
                                    file_name=file.name,
                                    mime="text/plain",
                                    key=f"download_{file.name}"
                                )
                            except Exception as e:
                                st.error(f"读取失败：{e}")
            else:
                st.info("工作目录为空，请先运行工作流")
        else:
            st.info("工作目录不存在")
    
    # 主界面
    display_agent_status_row()
    st.divider()
    
    display_active_agent_banner()
    
    display_chat_flow()
    
    # 用户输入
    st.divider()
    user_input = st.chat_input("💡 输入需求，或点击启动按钮开始...")
    
    if user_input:
        add_message("User_Proxy", user_input, "提出需求")
        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_msgs = [m for m in st.session_state.messages if m.get("agent_key") == "User_Proxy"]
        if not user_msgs:
            add_message("User_Proxy", "请开发一个电池健康状态预测 API", "提出需求")
        
        latest_requirement = user_msgs[-1]["content"] if user_msgs else "请开发一个电池健康状态预测 API"
        
        st.session_state.is_running = True
        
        # 进度提示
        progress_placeholder = st.empty()
        progress_placeholder.info("🚀 启动 SDLC 工作流，Agent 开始协作...")
        
        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
            )
            
            # 添加 Orchestrator 启动消息
            add_message("Orchestrator", f"🚀 启动 SDLC 工作流！用户需求：{latest_requirement[:100]}...", "启动流程")
            
            # 创建 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 → 汇总报告

开始协作！每个步骤完成后请明确说明。
"""
            
            # 执行对话
            with st.spinner("💬 Agent 们正在协作中，请稍候..."):
                chat_result = user_proxy.initiate_chat(
                    manager,
                    message=initial_message,
                    max_turns=max_round
                )
            
            # 记录所有对话
            for msg in groupchat.messages:
                agent_name = msg.get("name", "Unknown")
                content = msg.get("content", "")
                
                # 推断任务
                task_map = {
                    "PM_Agent": "需求分析",
                    "QA_Agent": "测试设计",
                    "Dev_Agent": "代码实现",
                    "Orchestrator": "流程协调",
                    "User_Proxy": "测试执行"
                }
                task = task_map.get(agent_name, "工作中")
                
                add_message(agent_name, content, task)
            
            # 完成
            add_message("Orchestrator", "✅ SDLC 流程完成！所有任务已完成。", "总结完成")
            
            # 保存生成的文件
            progress_placeholder.info("💾 正在保存生成的文件...")
            saved_files = save_generated_files()
            
            if saved_files:
                progress_placeholder.success(f"✅ SDLC 工作流完成！已保存 {len(saved_files)} 个文件到 workspace/ 目录")
            else:
                progress_placeholder.success("✅ SDLC 工作流完成！查看上方的对话流了解详情。")
            
            st.session_state.is_running = False
            
            st.rerun()
            
        except Exception as e:
            st.session_state.is_running = False
            st.session_state.current_agent = None
            progress_placeholder.error(f"❌ 错误：{str(e)}")
            st.error("请检查 API Key 和网络连接")
    
    if stop_btn:
        st.session_state.is_running = False
        st.session_state.current_agent = None
        st.info("⏸️ 工作流已暂停")


if __name__ == "__main__":
    main()