"""
AI 博客助手 - 多Agent协作系统
使用 LangGraph + Streamlit + OpenAI/Claude 实现

功能：
1. 研究员 Agent：搜索和收集信息
2. 作家 Agent：撰写博客内容
3. 编辑 Agent：审核和优化内容
"""

import streamlit as st
from typing import TypedDict, Annotated, List
import operator
from datetime import datetime
import json
import os

# 检查并安装依赖
try:
    from langgraph.graph import StateGraph, END
    from langchain_core.messages import HumanMessage, SystemMessage
except ImportError:
    st.error("""
    ⚠️ 缺少必要的依赖包！请运行：
    ```
    pip install langgraph langchain langchain-core langchain-openai
    ```
    """)
    st.stop()

# ============= 状态定义 =============
class BlogState(TypedDict):
    topic: str  # 博客主题
    research_notes: str  # 研究笔记
    draft: str  # 草稿
    final_blog: str  # 最终博客
    feedback: str  # 编辑反馈
    messages: Annotated[List[str], operator.add]  # 消息历史
    revision_count: int  # 修订次数
    use_mock: bool  # 是否使用模拟模式


# ============= LLM 配置 =============
def get_llm_response(prompt: str, system_prompt: str, use_mock: bool = False):
    """获取 LLM 响应（支持 OpenAI 或模拟模式）"""
    
    if use_mock:
        # 模拟模式：返回预设响应
        return f"[模拟响应] 基于提示生成的内容：\n{prompt[:100]}..."
    
    try:
        # 尝试使用 OpenAI
        from langchain_openai import ChatOpenAI
        
        api_key = os.environ.get("OPENAI_API_KEY") or st.session_state.get("openai_key")
        
        if not api_key:
            return "[错误] 请配置 OpenAI API Key"
        
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0.7,
            api_key=api_key
        )
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=prompt)
        ]
        
        response = llm.invoke(messages)
        return response.content
        
    except Exception as e:
        return f"[LLM 调用失败] {str(e)}\n\n使用模拟模式生成内容..."


# ============= Agent 节点 =============
class ResearcherAgent:
    """研究员 Agent - 负责收集信息"""
    
    def __call__(self, state: BlogState) -> BlogState:
        topic = state["topic"]
        use_mock = state.get("use_mock", True)
        
        system_prompt = """你是一位专业的研究员，擅长收集和整理信息。
请为给定的博客主题提供详细的研究笔记，包括：
1. 主题背景和重要性
2. 关键概念和术语
3. 目标读者分析
4. 建议的写作角度
5. 相关案例和数据"""

        user_prompt = f"""请为以下博客主题进行研究：
主题：{topic}

请提供结构化的研究笔记。"""

        if use_mock:
            # 模拟响应
            research = f"""# {topic} 研究笔记

## 1. 主题背景
{topic}是当前技术领域的重要话题，具有广泛的应用价值和发展前景。

## 2. 关键概念
- 核心概念：{topic}的基本定义和原理
- 技术栈：相关技术和工具
- 应用场景：实际使用案例

## 3. 目标读者
- 初学者：需要入门指导
- 进阶开发者：寻求最佳实践
- 技术决策者：关注价值和ROI

## 4. 建议角度
- 实用性：提供可操作的指南
- 深度：探讨技术细节
- 前瞻性：展望未来趋势

## 5. 参考资料
- 官方文档
- 技术博客
- 开源项目"""
        else:
            research = get_llm_response(user_prompt, system_prompt, use_mock)
        
        state["research_notes"] = research
        state["messages"].append(f"✅ 研究员完成调研：已收集 {topic} 相关信息")
        
        return state


class WriterAgent:
    """作家 Agent - 负责撰写内容"""
    
    def __call__(self, state: BlogState) -> BlogState:
        topic = state["topic"]
        research = state["research_notes"]
        feedback = state.get("feedback", "")
        use_mock = state.get("use_mock", True)
        
        if feedback:
            system_prompt = """你是一位经验丰富的技术博客作家。
请根据编辑的反馈修改博客内容，确保：
1. 解决所有提出的问题
2. 保持专业和易读
3. 添加必要的示例和说明"""

            user_prompt = f"""请修改以下博客内容：

原始主题：{topic}

研究笔记：
{research}

编辑反馈：
{feedback}

请提供修订后的完整博客内容。"""
        else:
            system_prompt = """你是一位经验丰富的技术博客作家。
请撰写高质量的博客文章，要求：
1. 结构清晰，有引言、正文、总结
2. 内容准确，有深度
3. 包含代码示例（如适用）
4. 语言流畅，易于理解
5. 使用 Markdown 格式"""

            user_prompt = f"""请根据以下研究笔记撰写博客：

主题：{topic}

研究笔记：
{research}

请撰写一篇完整的博客文章。"""

        if use_mock:
            # 模拟响应
            draft = f"""# {topic}：深入解析与实践指南

*作者：AI博客助手 | 日期：{datetime.now().strftime('%Y-%m-%d')}*

## 引言

在当今快速发展的技术领域，{topic}正在成为开发者和企业关注的焦点。本文将全面介绍{topic}的核心概念、应用场景和实践经验，帮助读者深入理解并掌握这一技术。

## 什么是{topic}？

{topic}是一种[技术描述]，它通过[工作原理]来实现[核心功能]。与传统方法相比，{topic}具有以下优势：

- **优势一**：提高效率和性能
- **优势二**：降低复杂度
- **优势三**：增强可维护性

## 核心特性

### 1. 特性一：创新性
{topic}采用了创新的方法来解决传统问题...

### 2. 特性二：可扩展性
系统架构设计灵活，支持水平和垂直扩展...

### 3. 特性三：易用性
提供友好的 API 和工具，降低使用门槛...

## 快速开始

让我们通过一个简单的例子来了解如何使用{topic}：

```python
# 示例代码
def example_function():
    \"\"\"这是一个{topic}的简单示例\"\"\"
    # 初始化
    config = {{
        'name': '{topic}',
        'version': '1.0'
    }}
    
    # 执行操作
    result = process(config)
    
    return result

# 运行示例
if __name__ == "__main__":
    output = example_function()
    print(f"结果: {{output}}")
```

## 实践案例

### 案例一：实际应用场景
在[场景描述]中，我们使用{topic}实现了[功能]...

### 案例二：性能优化
通过应用{topic}，系统性能提升了[数据]...

## 最佳实践

基于实际经验，以下是使用{topic}的最佳实践：

1. **从简单开始**：先掌握基础功能，再深入高级特性
2. **阅读文档**：官方文档是最好的学习资源
3. **参与社区**：加入技术社区，交流经验
4. **持续学习**：关注最新发展和更新

## 常见问题

**Q1: {topic}适合什么场景？**
A: {topic}特别适合[场景列表]...

**Q2: 如何优化性能？**
A: 可以通过[优化方法]来提升性能...

**Q3: 有哪些注意事项？**
A: 需要注意[注意事项列表]...

## 未来展望

{topic}的发展前景广阔，未来可能会看到：
- 更多的集成和生态系统
- 性能和功能的持续优化
- 更广泛的行业应用

## 总结

通过本文，我们全面了解了{topic}的核心概念、应用实践和最佳经验。{topic}作为一项重要技术，值得每位开发者学习和掌握。

希望这篇文章对你有所帮助。如果有任何问题或建议，欢迎在评论区讨论！

---

**参考资料**
- 官方文档
- 技术社区
- 开源项目

**标签**: {topic}, 技术, 教程, 最佳实践

{f"*本文修订次数: {state['revision_count']}*" if state['revision_count'] > 0 else ""}
"""
        else:
            draft = get_llm_response(user_prompt, system_prompt, use_mock)
        
        state["draft"] = draft
        state["messages"].append(f"✍️ 作家完成{'修订' if feedback else '初稿'}撰写")
        
        return state


class EditorAgent:
    """编辑 Agent - 负责审核和优化"""
    
    def __call__(self, state: BlogState) -> BlogState:
        draft = state["draft"]
        revision_count = state.get("revision_count", 0)
        use_mock = state.get("use_mock", True)
        
        # 质量检查
        issues = []
        
        if len(draft) < 500:
            issues.append("内容长度不足，需要扩充至至少500字")
        
        if "```" not in draft and ("代码" in state["topic"] or "编程" in state["topic"]):
            issues.append("技术文章缺少代码示例")
        
        if "总结" not in draft and "结论" not in draft:
            issues.append("缺少总结或结论部分")
            
        if draft.count("#") < 3:
            issues.append("文章结构层次不够清晰，建议增加小节")
        
        # 决定是否需要修订
        if issues and revision_count < 2:
            feedback_text = f"""编辑审核反馈（第{revision_count + 1}次）：

需要改进的地方：
{chr(10).join(f'{i+1}. {issue}' for i, issue in enumerate(issues))}

请针对以上问题进行修改，提升文章质量。"""
            
            state["feedback"] = feedback_text
            state["revision_count"] = revision_count + 1
            state["messages"].append(f"📝 编辑提出修改意见（第{revision_count + 1}次）")
            return state
        else:
            # 批准发布
            state["final_blog"] = draft
            state["feedback"] = ""
            
            if issues:
                state["messages"].append(f"✅ 编辑批准发布（达到最大修订次数，存在{len(issues)}个小问题）")
            else:
                state["messages"].append("✅ 编辑批准发布（质量优秀）")
            
            return state


# ============= 路由逻辑 =============
def should_continue(state: BlogState) -> str:
    """决定工作流下一步"""
    if state.get("final_blog"):
        return "end"
    elif state.get("feedback"):
        return "revise"
    else:
        return "review"


# ============= 构建工作流 =============
@st.cache_resource
def create_blog_workflow():
    """创建博客生成工作流（缓存）"""
    
    workflow = StateGraph(BlogState)
    
    # 添加节点
    workflow.add_node("researcher", ResearcherAgent())
    workflow.add_node("writer", WriterAgent())
    workflow.add_node("editor", EditorAgent())
    
    # 定义流程
    workflow.set_entry_point("researcher")
    workflow.add_edge("researcher", "writer")
    workflow.add_edge("writer", "editor")
    
    # 条件路由
    workflow.add_conditional_edges(
        "editor",
        should_continue,
        {
            "end": END,
            "revise": "writer",
            "review": "editor"
        }
    )
    
    return workflow.compile()


# ============= Streamlit 界面 =============
def main():
    st.set_page_config(
        page_title="AI 博客助手",
        page_icon="✍️",
        layout="wide"
    )
    
    # 初始化 session state
    if "openai_key" not in st.session_state:
        st.session_state.openai_key = ""
    
    st.title("✍️ AI 博客助手")
    st.markdown("### 🤖 多 Agent 协作的智能博客生成系统")
    
    # 侧边栏配置
    with st.sidebar:
        st.header("⚙️ 系统配置")
        
        # API 配置
        with st.expander("🔑 API 配置", expanded=False):
            api_key = st.text_input(
                "OpenAI API Key",
                type="password",
                value=st.session_state.openai_key,
                help="如果不提供，将使用模拟模式"
            )
            st.session_state.openai_key = api_key
            
            if api_key:
                st.success("✅ API Key 已配置")
            else:
                st.info("💡 未配置 API Key，将使用模拟模式")
        
        st.divider()
        
        # 系统架构说明
        st.markdown("""
        ### 🏗️ 系统架构
        
        **1. 研究员 Agent** 🔍
        - 收集主题相关信息
        - 分析目标受众
        - 提供写作建议
        
        **2. 作家 Agent** ✍️
        - 撰写博客内容
        - 根据反馈修订
        - 保持风格一致
        
        **3. 编辑 Agent** 📝
        - 审核内容质量
        - 检查结构完整
        - 提供改进建议
        """)
        
        st.divider()
        
        max_revisions = st.slider(
            "最大修订次数",
            min_value=0,
            max_value=5,
            value=2,
            help="编辑最多要求修订的次数"
        )
        
        st.divider()
        
        # 示例主题
        st.markdown("""
        ### 💡 示例主题
        - Python 异步编程
        - Docker 容器化
        - 微服务架构
        - React Hooks
        - 机器学习入门
        """)
    
    # 主界面
    col1, col2 = st.columns([3, 2])
    
    with col1:
        st.subheader("📝 输入博客主题")
        topic = st.text_input(
            "请输入主题",
            placeholder="例如：Python装饰器详解、Kubernetes入门指南等",
            help="输入您想要撰写的博客主题",
            label_visibility="collapsed"
        )

        # ✅ 多选版本
        quick_topics = st.multiselect(
            "快速选择（可多选）：",
            ["Python异步编程", "React Hooks实战", "Docker容器化", "机器学习基础", "API设计"]
        )

        
        if quick_topics:
            topic = quick_topics
    
    with col2:
        st.subheader("📊 工作流程")
        st.code("""
    研究员 → 作家 → 编辑
       ↓       ↓      ↓
     调研   撰写   审核
               ↺ 修订循环
        """, language="")
    
    # 生成按钮
    generate_btn = st.button(
        "🚀 开始生成博客",
        type="primary",
        use_container_width=True,
        disabled=not topic
    )
    
    # 生成博客
    if generate_btn and topic:
        # 确定是否使用模拟模式
        use_mock = not bool(st.session_state.openai_key)
        
        if use_mock:
            st.info("💡 使用模拟模式生成（未配置 API Key）")
        
        # 初始化状态
        initial_state = BlogState(
            topic=topic,
            research_notes="",
            draft="",
            final_blog="",
            feedback="",
            messages=[],
            revision_count=0,
            use_mock=use_mock
        )
        
        # 创建容器
        progress_container = st.container()
        message_container = st.container()
        result_container = st.container()
        
        with progress_container:
            progress_bar = st.progress(0, text="初始化...")
            status_placeholder = st.empty()
        
        try:
            # 创建工作流
            app = create_blog_workflow()
            
            # 执行工作流
            result = None
            step_count = 0
            total_steps = 6  # 估计步骤数
            
            with message_container:
                st.subheader("📋 执行日志")
                log_placeholder = st.empty()
                logs = []
            
            for step_result in app.stream(initial_state):
                step_count += 1
                progress = min(int((step_count / total_steps) * 100), 95)
                progress_bar.progress(progress, text=f"执行中... {progress}%")
                
                # 获取当前状态
                current_state = list(step_result.values())[0]
                
                # 显示消息
                if "messages" in current_state:
                    new_messages = current_state["messages"]
                    if new_messages:
                        logs.extend(new_messages)
                        log_placeholder.text_area(
                            "日志",
                            "\n".join(f"• {msg}" for msg in logs),
                            height=200,
                            label_visibility="collapsed"
                        )
                
                result = current_state
            
            # 完成
            progress_bar.progress(100, text="✅ 完成！")
            status_placeholder.success("博客生成完成！")
            
            # 显示结果
            with result_container:
                st.divider()
                
                # 统计信息
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("修订次数", result.get("revision_count", 0))
                with col2:
                    st.metric("内容长度", f"{len(result.get('final_blog', ''))} 字符")
                with col3:
                    st.metric("执行步骤", step_count)
                
                # 过程详情
                with st.expander("📋 研究笔记", expanded=False):
                    st.markdown(result.get("research_notes", ""))
                
                if result.get("draft") != result.get("final_blog"):
                    with st.expander("📄 草稿历史", expanded=False):
                        st.markdown(result.get("draft", ""))
                
                # 最终博客
                st.divider()
                st.subheader("📰 最终博客")
                
                final_blog = result.get("final_blog", "")
                
                if final_blog:
                    # 显示博客
                    st.markdown(final_blog)
                    
                    st.divider()
                    
                    # 操作按钮
                    col1, col2, col3, col4 = st.columns(4)


                    if isinstance(topic, list):
                        topic_str = "_".join(str(x).strip().replace(' ', '_') for x in topic if x)
                    else:
                        topic_str = str(topic)

                    with col1:
                        st.download_button(
                            label="📥 下载 MD",
                            data=final_blog,
                            file_name=f"{topic_str.replace(' ', '_')}.md",
                            mime="text/markdown",
                            use_container_width=True
                        )
                    
                    with col2:
                        export_data = {
                            "topic": topic,
                            "content": final_blog,
                            "metadata": {
                                "revision_count": result.get("revision_count", 0),
                                "created_at": datetime.now().isoformat(),
                                "mode": "mock" if use_mock else "llm"
                            }
                        }
                        st.download_button(
                            label="📥 下载 JSON",
                            data=json.dumps(export_data, ensure_ascii=False, indent=2),
                            file_name=f"{topic_str.replace(' ', '_')}.json",
                            mime="application/json",
                            use_container_width=True
                        )
                    
                    with col3:
                        if st.button("📋 复制内容", use_container_width=True):
                            st.code(final_blog, language="markdown")
                    
                    with col4:
                        if st.button("🔄 重新生成", use_container_width=True):
                            st.rerun()
                else:
                    st.error("生成失败，请重试")
        
        except Exception as e:
            st.error(f"❌ 生成过程出错: {str(e)}")
            with st.expander("查看错误详情"):
                st.exception(e)
    
    elif generate_btn:
        st.warning("⚠️ 请先输入博客主题")
    
    # 底部信息
    st.divider()
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("""
        ### 💡 使用提示
        - 输入明确具体的主题
        - 配置 API Key 获得更好效果
        - 查看执行日志了解过程
        - 支持多种格式下载
        """)
    
    with col2:
        st.markdown("""
        ### 🔧 技术栈
        - **LangGraph**: Agent 协作框架
        - **Streamlit**: 交互界面
        - **OpenAI**: 大语言模型
        - **Python**: 核心开发语言
        """)


if __name__ == "__main__":
    main()