"""
Streamlit 前端 - 多智能体实时聊天界面
提供可视化的 Agent 协同工作展示和人机交互
"""
import streamlit as st
import os
import sys
from pathlib import Path
from datetime import datetime
import time
from typing import Dict, Any, List

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

from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager
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-message {
    padding: 1rem;
    border-radius: 0.5rem;
    margin-bottom: 1rem;
}
.pm-agent { background-color: #e3f2fd; border-left: 4px solid #2196f3; }
.qa-agent { background-color: #e8f5e9; border-left: 4px solid #4caf50; }
.dev-agent { background-color: #fff3e0; border-left: 4px solid #ff9800; }
.orchestrator { background-color: #f3e5f5; border-left: 4px solid #9c27b0; }
.system { background-color: #fce4ec; border-left: 4px solid #e91e63; }
.user { background-color: #efebe9; border-left: 4px solid #795548; }
</style>
""", unsafe_allow_html=True)


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


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 get_agent_color(agent_name: str) -> str:
    """根据 Agent 名称返回颜色类"""
    color_map = {
        "PM_Agent": "pm-agent",
        "QA_Agent": "qa-agent",
        "Dev_Agent": "dev-agent",
        "Orchestrator": "orchestrator",
        "User_Proxy": "user",
        "system": "system"
    }
    return color_map.get(agent_name, "system")


def display_message(agent_name: str, message: str, timestamp: str):
    """显示单条消息"""
    color_class = get_agent_color(agent_name)
    
    # 头像映射
    avatar_map = {
        "PM_Agent": "📋",
        "QA_Agent": "✅",
        "Dev_Agent": "💻",
        "Orchestrator": "🎯",
        "User_Proxy": "👤",
        "system": "⚙️"
    }
    
    avatar = avatar_map.get(agent_name, "🤖")
    
    with st.chat_message(name=agent_name.lower().replace("_", ""), avatar=avatar):
        st.markdown(f"**{avatar} {agent_name}** *({timestamp})*")
        st.markdown(f"<div class='agent-message {color_class}'>{message}</div>", 
                   unsafe_allow_html=True)


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", "")
        
        md_content += f"### [{timestamp}] {agent}\n\n"
        md_content += f"{content}\n\n"
        md_content += "---\n\n"
    
    return md_content


def main():
    """主应用"""
    init_session_state()
    
    # ===== 左侧边栏 - 配置区域 =====
    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=15,
            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)
        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:
                import json
                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"
                )
    
    # ===== 主区域 - 聊天展示 =====
    st.title("🤖 AutoGen SDLC 多智能体协同系统")
    st.markdown("**基于 AutoGen + Qwen3.5-flash 的端到端软件交付系统**")
    
    # 状态指示器
    status_col1, status_col2, status_col3 = st.columns(3)
    with status_col1:
        st.metric("对话轮数", len(st.session_state.messages))
    with status_col2:
        status = "🟢 运行中" if st.session_state.is_running else "🔴 已停止"
        st.metric("系统状态", status)
    with status_col3:
        result_status = "✅ 成功" if st.session_state.workflow_result and st.session_state.workflow_result.get("success") else "⏳ 未开始"
        st.metric("工作流状态", result_status)
    
    st.divider()
    
    # 聊天历史展示区
    chat_container = st.container()
    with chat_container:
        if not st.session_state.messages:
            st.info("👈 请在左侧配置后，输入需求并点击启动按钮开始工作流")
        else:
            for msg in st.session_state.messages:
                display_message(
                    agent_name=msg.get("agent_name", "Unknown"),
                    message=msg.get("message", ""),
                    timestamp=msg.get("timestamp", "")
                )
    
    # ===== 右侧边栏 - 进度和日志 =====
    with st.sidebar:
        st.divider()
        
        # 工作流进度
        st.subheader("📊 工作流进度")
        
        workflow_steps = [
            "需求分析",
            "测试设计",
            "代码实现",
            "测试执行",
            "最终验证"
        ]
        
        current_step = st.session_state.get("current_step", 0)
        for i, step in enumerate(workflow_steps):
            if i < current_step:
                st.success(f"✅ {step}")
            elif i == current_step:
                st.info(f"🔄 {step}")
            else:
                st.write(f"⚪ {step}")
    
    # ===== 用户输入区域 =====
    st.divider()
    
    # 用户需求输入
    user_input = st.chat_input("请输入您的软件需求...")
    
    if user_input:
        # 添加用户消息到历史
        user_msg = {
            "timestamp": datetime.now().isoformat(),
            "agent_name": "User",
            "role": "user",
            "message": user_input
        }
        st.session_state.messages.append(user_msg)
        st.session_state.conversation_history.append(user_msg)
        
        # 如果正在运行，添加到队列等待处理
        if st.session_state.is_running:
            # 这里会触发工作流执行
            pass
    
    # ===== 工作流执行逻辑 =====
    if start_btn:
        if not api_key:
            st.error("请先设置 API Key！")
            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
        
        # 创建占位符用于实时更新
        status_placeholder = st.empty()
        message_placeholder = st.empty()
        
        with status_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.current_step = 0
            
            # 启动对话（简化版本，实际应该用异步）
            with message_placeholder:
                st.info("💬 Agent 们正在协作中，请稍候...")
            
            # 这里简化处理，实际应该使用异步回调
            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", "")
                }
                st.session_state.messages.append(chat_msg)
            
            # 更新状态
            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.rerun()
            
        except Exception as e:
            st.session_state.is_running = False
            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.info("⏸️ 工作流已暂停")


if __name__ == "__main__":
    main()