#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版多Agent数据分析系统
集成所有增强版Agent和新功能

系统包含以下Agent:
1. Enhanced RAG Agent - 智能电信数据分析问答
2. Enhanced Python Agent - 高级Python代码执行和数据科学
3. Enhanced AMAP Agent - 全功能地理位置服务
4. Safe File Agent - 安全文件管理
5. SQL Agent - 数据库查询和分析
6. Terminal Command Agent - 安全终端命令执行
7. Web Search Agent - 智能网络搜索检索
8. Supervisor - 智能协调调度
"""

from __future__ import annotations
import os
import asyncio
from typing import Literal
from dotenv import load_dotenv 
load_dotenv(override=True)

from langchain_deepseek import ChatDeepSeek
from langgraph_supervisor import create_supervisor
from langchain.chat_models import init_chat_model

# 导入所有增强版Agent
from enhanced_rag_agent import enhanced_rag_agent
from enhanced_python_agent import enhanced_python_agent
from enhanced_amap_agent import enhanced_amap_agent
from safe_file_agent import safe_file_agent
from sql_agent import sql_agent
from terminal_command_agent import terminal_command_agent
from web_search_agent import web_search_agent

# 初始化模型
model = ChatDeepSeek(model="deepseek-chat")

# =============================================================================
# 增强版Supervisor
# =============================================================================

ENHANCED_SUPERVISOR_PROMPT = """你是一个顶级的AI项目总监，负责管理一个由七位专家组成的AI团队。
你的职责是根据用户的需求，智能地将任务分配给最合适的专家处理。

**🎯 你的团队成员档案：**

1. **`enhanced_rag_agent` (智能电信数据分析专家)**
   - **能力**: 深度理解电信客户流失分析，具备多轮对话记忆和智能问答能力
   - **调用时机**:
     - 电信行业客户流失相关的专业问题
     - 需要专业术语解释和业务洞察
     - 数据科学建模理论和方法咨询
     - 需要上下文关联分析的复杂问题

2. **`enhanced_python_agent` (高级Python数据科学专家)**
   - **能力**: 强大的Python代码执行、数据分析、机器学习和高级可视化
   - **调用时机**:
     - Python代码编写和执行
     - 数据科学分析和统计计算
     - 机器学习模型构建和评估
     - 高级数据可视化和图表生成
     - 数据处理和清洗任务

3. **`enhanced_amap_agent` (全功能地理位置服务专家)**
   - **能力**: 完整的地理位置服务，包括天气、导航、区域分析和行程规划
   - **调用时机**:
     - 天气查询和预报分析
     - 地理编码和位置转换
     - POI搜索和区域设施分析
     - 路径规划和交通分析
     - 批量地理数据处理

4. **`safe_file_agent` (安全文件管理专家)**
   - **能力**: 在安全沙盒环境中进行文件操作，包括读写、管理和组织
   - **调用时机**:
     - 文件创建、读取、修改、删除
     - 目录管理和文件组织
     - 文档编辑和内容处理
     - 配置文件管理
     - 日志文件分析

5. **`sql_agent` (数据库操作专家)**
   - **能力**: 数据库查询、数据提取和数据库管理
   - **调用时机**:
     - SQL查询和数据库操作
     - 数据提取和导出
     - 数据库结构分析
     - 为其他Agent准备数据

6. **`terminal_command_agent` (终端命令执行和文件操作专家)**
   - **能力**: 安全的终端命令执行、完整的文件系统操作和系统管理
   - **调用时机**:
     - 终端命令执行和系统操作
     - 文件的增删改查操作
     - 目录管理和文件组织
     - 系统信息查询和监控
     - 开发工具调用和脚本执行
     - 网络工具使用和数据下载

7. **`web_search_agent` (智能网络搜索检索专家)**
   - **能力**: 基于Tavily API的实时网络搜索、新闻检索、金融信息查询和URL内容提取
   - **调用时机**:
     - 需要最新网络信息和实时资讯
     - 新闻事件查询和时事分析
     - 金融市场信息和股票数据检索
     - 技术文档和学术资料搜索
     - 从指定URL提取和分析内容
     - 竞品分析和市场调研
     - 验证信息的准确性和时效性

**🔄 工作流程原则：**

- **智能路由**: 准确识别用户需求，选择最适合的专家
- **协作配合**: 当任务需要多个专家时，合理安排协作顺序
- **依赖管理**: 确保数据依赖关系得到满足（如先用sql_agent提取数据，再用enhanced_python_agent分析）
- **结果整合**: 将多个专家的结果有机整合，提供完整解决方案
- **用户体验**: 保持流畅的交互体验，及时反馈处理进度

**📋 决策过程：**

1. **需求分析**: 仔细分析用户请求的类型和复杂度
2. **专家选择**: 选择最匹配的专家或专家组合
3. **任务分解**: 将复杂任务分解为专家可处理的子任务
4. **执行监控**: 跟踪任务执行进度和质量
5. **结果验证**: 确保输出满足用户需求
6. **完成确认**: 当所有任务完成时，输出 `FINISH`

**🎯 可选择的下一步行动：**
`['enhanced_rag_agent', 'enhanced_python_agent', 'enhanced_amap_agent', 'safe_file_agent', 'sql_agent', 'terminal_command_agent', 'web_search_agent', 'FINISH']`

请根据用户需求，选择最合适的专家来处理任务！
"""

# 创建增强版Supervisor
enhanced_supervisor = create_supervisor(
    model=init_chat_model(model="deepseek-chat", model_provider="deepseek", temperature=0),
    agents=[enhanced_rag_agent, enhanced_python_agent, enhanced_amap_agent, safe_file_agent, sql_agent, terminal_command_agent, web_search_agent],
    prompt=ENHANCED_SUPERVISOR_PROMPT,
    add_handoff_back_messages=True
).compile()

# =============================================================================
# 系统信息和帮助功能
# =============================================================================

def get_system_info() -> str:
    """获取增强版系统信息"""
    info = """
🚀 增强版多Agent数据分析系统
=====================================

本系统集成了7个专业Agent，提供全方位的数据分析和服务能力：

📊 **Enhanced RAG Agent** - 智能电信数据分析专家
   • 深度电信行业知识问答
   • 多轮对话记忆能力
   • 智能查询优化和扩展
   • 专业术语解释和建议

🐍 **Enhanced Python Agent** - 高级Python数据科学专家
   • 强大的代码执行能力
   • 高级数据可视化（多种图表类型）
   • 机器学习建模和评估
   • 统计分析和数据探索

🗺️ **Enhanced AMAP Agent** - 全功能地理位置服务专家
   • 天气查询和预报分析
   • 智能行程规划
   • 区域设施分析
   • 交通状况分析
   • 批量地理数据处理

🔒 **Safe File Agent** - 安全文件管理专家
   • 沙盒环境文件操作
   • 安全的文件读写管理
   • 目录结构组织
   • 操作日志记录

💾 **SQL Agent** - 数据库操作专家
   • MySQL数据库查询
   • 数据提取和分析
   • 为其他Agent提供数据支持

🖥️ **Terminal Command Agent** - 终端命令执行和文件操作专家
   • 安全的终端命令执行
   • 完整的文件系统操作（增删改查）
   • 系统信息查询和监控
   • 开发工具调用和脚本执行
   • 多重安全防护机制

🔍 **Web Search Agent** - 智能网络搜索检索专家
   • 实时网络信息搜索（基于Tavily API）
   • 专业新闻和金融信息检索
   • URL内容智能提取和分析
   • 多主题搜索（通用、新闻、金融）
   • AI优化的搜索结果和摘要

🧠 **Intelligent Supervisor** - 智能协调调度
   • 自动识别用户需求
   • 智能选择合适的专家
   • 管理Agent间的协作
   • 整合多专家结果

=====================================
系统特性：
✅ 多Agent智能协作
✅ 安全性和稳定性保障
✅ 丰富的功能集成
✅ 用户友好的交互体验
✅ 全面的错误处理和日志
=====================================
"""
    return info

def get_usage_examples() -> str:
    """获取使用示例"""
    examples = """
💡 使用示例
=====================================

📊 **数据分析类问题**:
"客户流失预测模型有哪些关键特征？如何进行特征工程？"
→ 系统会调用 enhanced_rag_agent 提供专业解答

🐍 **Python编程和数据科学**:
"创建一个包含100个随机数的数据集，然后绘制直方图和箱线图"
→ 系统会调用 enhanced_python_agent 执行代码并生成可视化

🗺️ **地理位置服务**:
"查询北京和上海的天气，然后规划从北京到上海的最佳路线"
→ 系统会调用 enhanced_amap_agent 处理地理相关需求

🔒 **文件管理**:
"创建一个名为'project_notes.md'的文件，写入今天的工作总结"
→ 系统会调用 safe_file_agent 安全地处理文件操作

🖥️ **终端命令执行和系统操作**:
"查看当前目录内容，然后创建一个新文件夹并下载指定URL的文件"
→ 系统会调用 terminal_command_agent 执行终端命令和文件操作

💾 **数据库操作**:
"查询数据库中客户表的结构，然后提取前100条客户数据进行分析"
→ 系统会先调用 sql_agent 提取数据，再调用其他Agent分析

🔄 **复合任务**:
"从数据库提取客户数据，进行流失预测建模，并生成可视化报告"
→ 系统会协调多个Agent协作完成复杂任务

=====================================
"""
    return examples

# =============================================================================
# 主要接口函数
# =============================================================================

async def enhanced_system_query(user_query: str) -> str:
    """
    增强版系统查询接口
    """
    try:
        result = await enhanced_supervisor.ainvoke({
            "messages": [{"role": "user", "content": user_query}]
        })
        
        return result["messages"][-1].content
        
    except Exception as e:
        return f"❌ 系统处理失败: {str(e)}"

# =============================================================================
# 测试和演示功能
# =============================================================================

async def run_system_demo():
    """运行系统演示"""
    print(get_system_info())
    print(get_usage_examples())
    
    demo_queries = [
        "什么是客户流失？电信行业为什么要关注这个问题？",
        "帮我查询北京的天气情况",
        "创建一个包含20个随机数的列表，计算平均值并绘制柱状图",
        "在工作区创建一个hello.txt文件，内容是'Hello Enhanced System!'",
        "搜索Python 3.12的最新特性和改进",
        "查询最近一周的AI技术新闻",
    ]
    
    print("🎭 系统功能演示")
    print("=" * 50)
    
    for i, query in enumerate(demo_queries, 1):
        print(f"\n📝 演示 {i}: {query}")
        print("-" * 40)
        
        try:
            response = await enhanced_system_query(query)
            print(f"🤖 系统回答: {response[:300]}...")
        except Exception as e:
            print(f"❌ 演示失败: {str(e)}")
        
        print("\n" + "="*50)
        await asyncio.sleep(2)

if __name__ == "__main__":
    print("🚀 启动增强版多Agent数据分析系统...")
    
    # 显示系统信息
    print(get_system_info())
    
    # 运行演示
    try:
        asyncio.run(run_system_demo())
    except KeyboardInterrupt:
        print("\n⏹️ 系统演示被中断")
    except Exception as e:
        print(f"\n❌ 系统启动失败: {str(e)}")
        import traceback
        traceback.print_exc()