"""
Multi-Agent Decision Workshop - 主应用
多 Agent 决策工作坊：通过多角色辩论帮助用户做出更好的决策
"""
import streamlit as st
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

from agents import get_all_agents, get_recommended_agents, AGENT_PROFILES
from orchestrator import DebateManager, DebateConfig
from orchestrator.research_manager import ResearchManager, ResearchConfig
from report import ReportGenerator
from utils import LLMClient
import config

# ==================== 页面配置 ====================
st.set_page_config(
    page_title="🎭 多 Agent 决策工作坊",
    page_icon="🎭",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ==================== 样式 ====================
st.markdown("""
<style>
    .agent-card {
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 0.5rem;
        border-left: 4px solid #4A90A4;
        background-color: #f8f9fa;
    }
    .speech-bubble {
        background-color: #f0f2f6;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 0.5rem 0;
    }
    .round-header {
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 0.5rem 1rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
    }
    .custom-agent-form {
        background-color: #e8f4f8;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 0.5rem 0;
    }
    .research-step {
        border-left: 3px solid #FF4B4B;
        padding-left: 10px;
        margin-bottom: 10px;
    }
</style>
""", unsafe_allow_html=True)

# ==================== 常量定义 ====================
# 从环境变量读取 API Key（隐藏在 .env 文件中）
DEFAULT_API_KEY = os.getenv("AIHUBMIX_API_KEY", "")

# 支持的模型列表
AVAILABLE_MODELS = {
    "gpt-4o": "GPT-4o (推荐)",
    "gpt-4o-mini": "GPT-4o Mini (快速)",
    "gpt-4-turbo": "GPT-4 Turbo",
    "gpt-3.5-turbo": "GPT-3.5 Turbo (经济)",
    "claude-3-5-sonnet-20241022": "Claude 3.5 Sonnet",
    "claude-3-opus-20240229": "Claude 3 Opus",
    "claude-3-haiku-20240307": "Claude 3 Haiku (快速)",
    "deepseek-chat": "DeepSeek Chat",
    "deepseek-coder": "DeepSeek Coder",
    "gemini-1.5-pro": "Gemini 1.5 Pro",
    "gemini-1.5-flash": "Gemini 1.5 Flash",
    "qwen-turbo": "通义千问 Turbo",
    "qwen-plus": "通义千问 Plus",
    "glm-4": "智谱 GLM-4",
    "moonshot-v1-8k": "Moonshot (月之暗面)",
}

# 决策类型
DECISION_TYPES = {
    "product": "产品方案",
    "business": "商业决策", 
    "tech": "技术选型",
    "personal": "个人规划"
}

# ==================== 初始化 Session State ====================
if "mode" not in st.session_state:
    st.session_state.mode = "Deep Research"

# Debate State
if "debate_started" not in st.session_state:
    st.session_state.debate_started = False
if "debate_finished" not in st.session_state:
    st.session_state.debate_finished = False
if "speeches" not in st.session_state:
    st.session_state.speeches = []
if "report" not in st.session_state:
    st.session_state.report = ""
if "custom_agents" not in st.session_state:
    st.session_state.custom_agents = {}

# Research State
if "research_plan" not in st.session_state:
    st.session_state.research_plan = ""
if "research_started" not in st.session_state:
    st.session_state.research_started = False
if "research_output" not in st.session_state:
    st.session_state.research_output = "" # Final report
if "research_steps_output" not in st.session_state:
    st.session_state.research_steps_output = [] # List of step results


# ==================== 侧边栏：配置 ====================
with st.sidebar:
    st.header("⚙️ 设置")
    
    # 全局 API Key 设置
    with st.expander("🔑 API Key 设置", expanded=True):
        use_custom_key = st.checkbox("使用自定义 API Key")
        if use_custom_key:
            api_key = st.text_input(
                "API Key",
                type="password",
                help="留空则使用环境变量中的 Key"
            )
        else:
            api_key = DEFAULT_API_KEY
            
    st.divider()
    
    # 模式选择
    mode = st.radio(
        "📊 选择模式",
        ["Deep Research", "Debate Workshop"],
        index=0 if st.session_state.mode == "Deep Research" else 1
    )
    st.session_state.mode = mode
    
    st.divider()

    if mode == "Deep Research":
        st.subheader("🧪 研究模型配置")
        
        # 3 个角色的模型配置
        roles_config = {}
        for role_key, role_info in config.RESEARCH_MODEL_ROLES.items():
            roles_config[role_key] = st.selectbox(
                f"{role_info['name']} ({role_info['description']})",
                options=list(AVAILABLE_MODELS.keys()),
                index=list(AVAILABLE_MODELS.keys()).index(role_info['default_model']) if role_info['default_model'] in AVAILABLE_MODELS else 0,
                key=f"model_{role_key}"
            )
            
    else: # Debate Workshop
        # 模型选择
        model = st.selectbox(
            "🤖 选择通用模型",
            options=list(AVAILABLE_MODELS.keys()),
            format_func=lambda x: AVAILABLE_MODELS[x],
            index=0,
            help="选择用于辩论的 AI 模型"
        )
        
        # 辩论配置
        max_rounds = st.slider(
            "🔄 辩论轮数",
            min_value=1,
            max_value=4,
            value=2,
            help="每轮所有 Agent 都会发言一次"
        )
        
        st.divider()
        
        # ==================== 自定义角色 (Debate Only) ====================
        st.subheader("✨ 自定义角色")
        
        with st.expander("➕ 添加新角色", expanded=False):
            new_agent_name = st.text_input("角色名称", placeholder="如：法务顾问", key="new_agent_name")
            new_agent_emoji = st.text_input("角色 Emoji", value="🎯", max_chars=2, key="new_agent_emoji")
            new_agent_perspective = st.text_input("视角定位", placeholder="如：法律合规视角", key="new_agent_perspective")
            new_agent_focus = st.text_input("关注点（逗号分隔）", placeholder="如：合规风险, 法律条款", key="new_agent_focus")
            new_agent_prompt = st.text_area("角色设定 Prompt", placeholder="描述这个角色的思考方式...", height=100, key="new_agent_prompt")
            
            if st.button("✅ 添加角色", use_container_width=True):
                if new_agent_name and new_agent_prompt:
                    agent_id = f"custom_{len(st.session_state.custom_agents)}"
                    st.session_state.custom_agents[agent_id] = {
                        "name": new_agent_name,
                        "emoji": new_agent_emoji,
                        "perspective": new_agent_perspective or "自定义视角",
                        "focus_areas": [f.strip() for f in new_agent_focus.split(",") if f.strip()],
                        "system_prompt": new_agent_prompt
                    }
                    st.success(f"已添加角色: {new_agent_emoji} {new_agent_name}")
                    st.rerun()
                else:
                    st.warning("请至少填写角色名称和 Prompt")
        
        # 显示已添加的自定义角色
        if st.session_state.custom_agents:
            st.markdown("**已添加的自定义角色：**")
            for agent_id, agent_info in list(st.session_state.custom_agents.items()):
                col1, col2 = st.columns([3, 1])
                with col1:
                    st.markdown(f"{agent_info['emoji']} {agent_info['name']}")
                with col2:
                    if st.button("🗑️", key=f"del_{agent_id}"):
                        del st.session_state.custom_agents[agent_id]
                        st.rerun()

# ==================== 主界面逻辑 ====================

if mode == "Deep Research":
    st.title("🧪 Deep Research Mode")
    st.markdown("*深度研究模式：规划 -> 研究 -> 报告*")
    
    # Input
    research_topic = st.text_area("研究主题", placeholder="请输入你想深入研究的主题...", height=100)
    research_context = st.text_area("补充背景 (可选)", placeholder="任何额外的背景信息...", height=80)
    
    generate_plan_btn = st.button("📝 生成研究计划", type="primary", disabled=not research_topic)
    
    if generate_plan_btn and research_topic:
        st.session_state.research_started = False
        st.session_state.research_output = ""
        st.session_state.research_steps_output = []
        
        manager = ResearchManager(api_key=api_key)
        config_obj = ResearchConfig(
            topic=research_topic,
            context=research_context,
            planner_model=roles_config['planner'],
            researcher_model=roles_config['researcher'],
            writer_model=roles_config['writer']
        )
        manager.create_agents(config_obj)
        
        with st.spinner("正在制定研究计划..."):
            plan_text = ""
            for chunk in manager.generate_plan(research_topic, research_context):
                plan_text += chunk
            st.session_state.research_plan = plan_text
            
    # Plan Review & Edit
    if st.session_state.research_plan:
        st.divider()
        st.subheader("📋 研究计划确认")
        
        edited_plan = st.text_area("请审查并编辑计划 (Markdown格式)", value=st.session_state.research_plan, height=300)
        st.session_state.research_plan = edited_plan
        
        start_research_btn = st.button("🚀 开始深度研究", type="primary")
        
        if start_research_btn:
            st.session_state.research_started = True
            st.session_state.research_steps_output = [] # Reset steps
            
            # Parse plan lines to get steps (simple heuristic: lines starting with - or 1.)
            steps = [line.strip() for line in edited_plan.split('\n') if line.strip().startswith(('-', '*', '1.', '2.', '3.', '4.', '5.'))]
            if not steps:
                steps = [edited_plan] # Fallback if no list format
            
            manager = ResearchManager(api_key=api_key)
            config_obj = ResearchConfig(
                topic=research_topic,
                context=research_context,
                planner_model=roles_config['planner'],
                researcher_model=roles_config['researcher'],
                writer_model=roles_config['writer']
            )
            manager.create_agents(config_obj)
            
            # Execute Steps
            previous_findings = ""
            st.divider()
            st.subheader("🔍 研究进行中...")
            
            step_progress = st.container()
            
            for i, step in enumerate(steps):
                with step_progress:
                    with st.status(f"正在研究: {step}", expanded=True):
                        findings_text = ""
                        placeholder = st.empty()
                        for chunk in manager.execute_step(step, previous_findings):
                            findings_text += chunk
                            placeholder.markdown(findings_text)
                        
                        st.session_state.research_steps_output.append(f"### {step}\n{findings_text}")
                        previous_findings += f"\n\nFinding for '{step}':\n{findings_text}"
            
            # Final Report
            st.divider()
            st.subheader("📄 最终报告生成中...")
            report_placeholder = st.empty()
            final_report = ""
            for chunk in manager.generate_report(research_topic, previous_findings):
                final_report += chunk
                report_placeholder.markdown(final_report)
                
            st.session_state.research_output = final_report
            st.success("✅ 研究完成")
            
    # Show Final Report if available
    if st.session_state.research_output:
        st.divider()
        st.subheader("📄 最终研究报告")
        st.markdown(st.session_state.research_output)
        st.download_button("📥 下载报告", st.session_state.research_output, "research_report.md")


elif mode == "Debate Workshop":
    # ==================== 原始 Debate UI 逻辑 ====================
    st.title("🎭 多 Agent 决策工作坊")
    st.markdown("*让多个 AI 角色从不同视角辩论，帮助你做出更全面的决策*")
    
    # ==================== 输入区域 ====================
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.subheader("📝 决策议题")
        
        # 决策类型选择
        decision_type = st.selectbox(
            "决策类型",
            options=list(DECISION_TYPES.keys()),
            format_func=lambda x: DECISION_TYPES[x],
            index=0
        )
        
        # 议题输入
        topic = st.text_area(
            "请描述你的决策议题",
            placeholder="例如：我们是否应该在 Q2 推出 AI 助手功能？\n\n或者：我应该接受这份新工作 offer 吗？",
            height=120
        )
        
        # 背景信息（可选）
        with st.expander("➕ 添加背景信息（可选）"):
            context = st.text_area(
                "背景信息",
                placeholder="提供更多上下文信息，如：\n- 当前状况\n- 已有的资源和限制\n- 相关数据和事实",
                height=100
            )
        context = context if 'context' in dir() else ""
    
    with col2:
        st.subheader("🎭 选择参与角色")
        
        # 获取推荐的角色
        recommended = get_recommended_agents(decision_type)
        all_agents = get_all_agents()
        
        # 预设角色选择
        st.markdown("**预设角色：**")
        selected_agents = []
        for agent in all_agents:
            is_recommended = agent["id"] in recommended
            default_checked = is_recommended
            
            if st.checkbox(
                f"{agent['emoji']} {agent['name']}",
                value=default_checked,
                key=f"agent_{agent['id']}"
            ):
                selected_agents.append(agent["id"])
        
        # 自定义角色选择
        if st.session_state.custom_agents:
            st.markdown("**自定义角色：**")
            for agent_id, agent_info in st.session_state.custom_agents.items():
                if st.checkbox(
                    f"{agent_info['emoji']} {agent_info['name']}",
                    value=True,
                    key=f"agent_{agent_id}"
                ):
                    selected_agents.append(agent_id)
        
        # 角色数量提示
        if len(selected_agents) < 2:
            st.warning("请至少选择 2 个角色")
        elif len(selected_agents) > 6:
            st.warning("建议不超过 6 个角色")
        else:
            st.info(f"已选择 {len(selected_agents)} 个角色")
    
    # ==================== 辩论控制 ====================
    st.divider()
    
    col_btn1, col_btn2, col_btn3 = st.columns([1, 1, 2])
    
    with col_btn1:
        start_btn = st.button(
            "🚀 开始辩论",
            disabled=(not topic or len(selected_agents) < 2 or not api_key),
            type="primary",
            use_container_width=True
        )
    
    with col_btn2:
        reset_btn = st.button(
            "🔄 重置",
            use_container_width=True
        )
    
    if reset_btn:
        st.session_state.debate_started = False
        st.session_state.debate_finished = False
        st.session_state.speeches = []
        st.session_state.report = ""
        st.rerun()
    
    # ==================== 辩论展示区 ====================
    if start_btn and topic and len(selected_agents) >= 2:
        st.session_state.debate_started = True
        st.session_state.speeches = []
        
        st.divider()
        st.subheader("🎬 辩论进行中...")
        
        # 临时将自定义角色添加到 agent_profiles
        from agents import agent_profiles
        original_profiles = dict(agent_profiles.AGENT_PROFILES)
        agent_profiles.AGENT_PROFILES.update(st.session_state.custom_agents)
        
        try:
            # 初始化客户端和管理器
            provider_val = "aihubmix" # Debate mode default to aihubmix or logic needs to be robust. 
            # Note: in sidebar "model" and "api_key" were set. "provider" variable is now inside the Sidebar logic block if mode==Debate.
            # But wait, I removed the "Advanced Settings" block from the global scope and put it into sub-scope?
            # Let's check my sidebar logic above.
            
            # Refactoring check:
            # I removed the provider selection logic from the global sidebar. I should probably add it back or assume a default.
            # In the original code, provider selection was in "Advanced Settings".
            
            llm_client = LLMClient(
                provider="aihubmix", 
                api_key=api_key,
                base_url="https://aihubmix.com/v1",
                model=model
            )
            debate_manager = DebateManager(llm_client)
            
            # 配置辩论
            debate_config = DebateConfig(
                topic=topic,
                context=context,
                agent_ids=selected_agents,
                max_rounds=max_rounds
            )
            debate_manager.setup_debate(debate_config)
            
            # 运行辩论（流式）
            current_round = 0
            speech_placeholder = None
            
            for event in debate_manager.run_debate_stream():
                if event["type"] == "round_start":
                    current_round = event["round"]
                    st.markdown(
                        f'<div class="round-header">📢 第 {current_round} 轮讨论</div>',
                        unsafe_allow_html=True
                    )
                
                elif event["type"] == "speech_start":
                    st.markdown(f"**{event['emoji']} {event['agent_name']}**")
                    speech_placeholder = st.empty()
                    current_content = ""
                
                elif event["type"] == "speech_chunk":
                    current_content += event["chunk"]
                    speech_placeholder.markdown(current_content)
                
                elif event["type"] == "speech_end":
                    st.session_state.speeches.append({
                        "agent_id": event["agent_id"],
                        "content": event["content"],
                        "round": current_round
                    })
                    st.divider()
                
                elif event["type"] == "debate_end":
                    st.session_state.debate_finished = True
                    st.success("✅ 辩论结束！正在生成决策报告...")
            
            # 生成报告
            if st.session_state.debate_finished:
                report_generator = ReportGenerator(llm_client)
                speeches = debate_manager.get_all_speeches()
                
                st.subheader("📊 决策报告")
                report_placeholder = st.empty()
                report_content = ""
                
                for chunk in report_generator.generate_report_stream(
                    topic=topic,
                    speeches=speeches,
                    context=context
                ):
                    report_content += chunk
                    report_placeholder.markdown(report_content)
                
                st.session_state.report = report_content
                
                # 下载按钮
                st.download_button(
                    label="📥 下载报告 (Markdown)",
                    data=report_content,
                    file_name="decision_report.md",
                    mime="text/markdown"
                )
        
        except Exception as e:
            st.error(f"发生错误: {str(e)}")
            import traceback
            st.code(traceback.format_exc())
            st.info("请检查你的 API Key 和模型设置是否正确")
        
        finally:
            # 恢复原始角色配置
            agent_profiles.AGENT_PROFILES = original_profiles
    
    # ==================== 历史报告展示 ====================
    elif st.session_state.report and not start_btn:
        st.divider()
        st.subheader("📊 上次的决策报告")
        st.markdown(st.session_state.report)
        
        st.download_button(
            label="📥 下载报告 (Markdown)",
            data=st.session_state.report,
            file_name="decision_report.md",
            mime="text/markdown"
        )

# ==================== 底部信息 ====================
st.divider()
col_footer1, col_footer2, col_footer3 = st.columns(3)
with col_footer2:
    st.markdown(
        "<div style='text-align: center; color: #888;'>"
        "🎭 Multi-Agent Decision Workshop<br>多 Agent 决策工作坊"
        "</div>",
        unsafe_allow_html=True
    )