import os
from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
from qwen_agent.tools.base import BaseTool, register_tool
import json5
import sys
from datetime import datetime
import pandas as pd
from docx import Document

# 获取当前脚本所在目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 获取当前运行代码的解释器路径
current_interpreter = sys.executable

# 获取当前时间
@register_tool('get_current_datetime')
class GetCurrentDateTime(BaseTool):
    """获取当前系统日期时间的工具，返回年、月、日、时、分、秒等信息"""
    description = '获取当前系统日期时间的工具，可返回完整的日期时间信息，用于辅助旅行规划中对时间的判断。'

    parameters = []  # 此工具不需要参数

    def __init__(self, *args, **kwargs):
        super().__init__()
        pass

    def call(self, params: str, **kwargs) -> str:
        try:
            # 获取当前时间
            now = datetime.now()

            # 构建响应数据
            current_time_info = {
                'year': now.year,
                'month': now.month,
                'day': now.day,
                'hour': now.hour,
                'minute': now.minute,
                'second': now.second,
                'formatted_time': now.strftime('%Y-%m-%d %H:%M:%S')  # 标准格式时间字符串
            }

            response = {
                'status': 'success',
                'current_datetime': current_time_info
            }

            return json5.dumps(response, ensure_ascii=False)

        except Exception as e:
            return json5.dumps({
                'status': 'error',
                'message': f'获取当前时间失败: {str(e)}'
            })

# 获取南航航班信息
@register_tool('flight_search')
class FlightSearch(BaseTool):
    """航班信息检索工具，根据出发地、目的地和日期查询航班信息"""
    description = '航班信息检索工具，根据出发地、目的地和日期查询航班信息，返回符合条件的航班列表。'
    parameters = [
        {
            'name': 'departure_city',
            'type': 'string',
            'description': '出发城市名称。例如：广州、北京',
            'required': True
        },
        {
            'name': 'arrival_city',
            'type': 'string',
            'description': '到达城市名称，例如：上海、杭州',
            'required': True
        },
        {
            'name': 'date',
            'type': 'string',
            'description': '查询日期，格式：YYYY/MM/DD，例如：2025/04/29',
            'required': True
        }
    ]

    def __init__(self, *args, **kwargs):
        """初始化工具，加载航班数据"""
        super().__init__()
        self.csv_file_path = os.path.join(BASE_DIR, "Flight.csv")
        self.data = self._load_data()

    def _load_data(self):
        """加载CSV数据并进行预处理"""
        try:
            df = pd.read_csv(self.csv_file_path)
            if '航班日期' in df.columns:
                # 将日期列转换为字符串格式 YYYY/MM/DD
                df['航班日期_str'] = pd.to_datetime(df['航班日期']).dt.strftime('%Y/%m/%d')
            return df
        except Exception as e:
            print(f"加载航班数据失败: {e}")
            return pd.DataFrame()

    def call(self, params: str, **kwargs) -> str:
        try:
            params = json5.loads(params)
            departure = params['departure_city']
            arrival = params['arrival_city']
            date_str = params['date']  # 保持字符串格式，如 '2025/5/30'

            if self.data.empty:
                return json5.dumps({
                   'status': 'error',
                   'message': '航班数据未成功加载，请检查CSV文件路径'
                })

            # 执行筛选查询（使用字符串列进行比较）
            filtered_df = self.data[
                (self.data['出发城市'] == departure) &
                (self.data['终到城市'] == arrival) &
                (self.data['航班日期_str'] == date_str)  # 关键修改点
            ]

            results = filtered_df.to_dict('records')

            response = {
               'status':'success',
                'query': {
                    'departure_city': departure,
                    'arrival_city': arrival,
                    'date': date_str  # 保持原始字符串格式
                },
                'total_results': len(results),
                'flights': results
            }

            return json5.dumps(response, ensure_ascii=False)
        except Exception as e:
            return json5.dumps({
               'status': 'error',
               'message': f'查询过程中发生错误: {str(e)}'
            })

@register_tool('get_service_info')
class GetServiceInfo(BaseTool):
    """获取南航服务信息的工具，根据服务名称读取对应docx文件内容"""
    description = '根据输入的服务名称（如“机上餐食”）读取对应的文档内容并返回'
    parameters = [
        {
            'name': 'service_name',
            'type': 'string',
            'description': '服务名称，可选值：机上餐食_公务舱、机上餐食_经济舱、机上餐食_特殊餐食、机上餐食_绿色飞行、机舱特色_公务舱、机舱特色_经济舱、机上娱乐、机上上网、机上免税品服务',
            'required': True,
            'enum': ['机上餐食_公务舱', '机上餐食_经济舱' , '机上餐食_特殊餐食', '机上餐食_绿色飞行', '机舱特色_公务舱', '机舱特色_经济舱', '机上娱乐', '机上上网','机上免税品服务']
        }
    ]

    def __init__(self, *args, **kwargs):
        super().__init__()
        self.docs_dir = os.path.join(BASE_DIR, "service_docx")

    def _read_docx(self, file_path):
        """读取.docx文件内容，返回纯文本"""
        try:
            doc = Document(file_path)
            full_text = []
            for para in doc.paragraphs:
                full_text.append(para.text)
            return '\n'.join(full_text)
        except Exception as e:
            return f"读取文档失败: {str(e)}"

    def call(self, params: str, **kwargs) -> str:
        try:
            params = json5.loads(params)
            service_name = params['service_name']

            # 构建文件路径
            file_name = f"{service_name}.docx"
            file_path = os.path.join(self.docs_dir, file_name)

            if not os.path.exists(file_path):
                return json5.dumps({
                    'status': 'error',
                    'message': f'未找到对应文档：{file_name}'
                }, ensure_ascii=False)

            content = self._read_docx(file_path)

            return json5.dumps({
                'status': 'success',
                'service_name': service_name,
                'content': content
            }, ensure_ascii=False)

        except Exception as e:
            return json5.dumps({
                'status': 'error',
                'message': f'查询过程中发生错误: {str(e)}'
            }, ensure_ascii=False)

def init_agent_service_railway_plan():
    llm_cfg = {
        # Use the model service provided by DashScope:
        'model': 'qwen3-235b-a22b',
        'model_server': 'dashscope',
        'api_key': "sk-1eb585dbfeea444691cd4216fd2a071d",
        'generate_cfg': {
            'top_p': 0.8
        }
    }
    tools = [
        # 查询高铁MCP server
        {
            "mcpServers": {
                "tongchenglvxing-mcp-server": {
                    "command": "npx",
                    "args": [
                        "-y",
                        "@wuchubuzai/tongchenglvxing-mcp-server"
                    ]
                }
            }
        },

        'get_current_datetime',
        'code_interpreter'
    ]
    system_message = """
    你是一个高铁规划助手，根据用户的输入，为其推荐相应的高铁班次。请满足如下要求：
    1) 如果用户并没有提供明确的时间信息，请使用get_current_datetime工具获取当前时间来推断用户要求时间。
    2) 请在回答的末尾加上信息来源，如“高铁信息来源：同程旅行”
    """
    bot = Assistant(llm=llm_cfg,
                    function_list=tools,
                    name='高铁推荐助手',
                    system_message=system_message,
                    description="我是高铁查询助手，可根据您的需求为您推荐高铁。")

    return bot

def init_agent_service_flight_plan():
    llm_cfg = {
        # Use the model service provided by DashScope:
        'model': 'qwen3-235b-a22b',
        'model_server': 'dashscope',
        'api_key': "sk-1eb585dbfeea444691cd4216fd2a071d",
        'generate_cfg': {
            'top_p': 0.8
        }
    }
    tools = [
        # 查询航班 MCP server
        {
            "mcpServers": {
                "variflight": {
                    "command": "npx",
                    "args": [
                        "-y",
                        "@variflight-ai/variflight-mcp"
                    ],
                    "env": {
                        "X_VARIFLIGHT_KEY": "sk-6Si9JTjgDINieKYLEE2VNMyZI285AuCYF8wl3Z-n18c"
                    }
                }
            }
        },
        'get_service_info',
        'flight_search',
        'get_current_datetime',
        'code_interpreter'
    ]
    system_message = """
    你是一个航班推荐助手，根据用户的输入，为其推荐相应的航班班次。请按照以下要求为用户推荐航班：
    1) 如果用户需要航班推荐，请优先使用flight_search工具推荐南方航空航班。
    2) 如果flight_search工具获取不到相关航班，则使用variflight MCP server推荐航班，同样优先推荐南航航班。
    3) 如果用户并没有提供明确的时间信息，请使用get_current_datetime工具获取当前时间来推断用户要求时间。
    4) 请在回答的末尾加上信息来源，如果航安信息是来自于工具flight_search则为“航班信息来源：中国南方航空公司”； 如果是来自于variflight则为“航班信息来源：Variflight航班平台”。
    
    同时你也是一位南航服务的宣传助手，在向用户推荐完南航航班之后，请你向用户主动宣传南航特色机上服务，包括机上餐食、机舱特色、机上娱乐、机上上网、机上免税品服务等。
    根据用户反馈，使用'get_service_info'工具获取具体服务信息再给用户介绍，采用图文的多模态展示和介绍。
    文末的宣传词如下(### ###内的内容)，可以根据情况适当更改和优化：
    ###
    ✈️ 中国南方航空 —— 让您的旅程更舒适、更精彩 ✈️

    选择南航，不仅是选择一次飞行，更是开启一段高品质的空中旅程！
    
    💺 宽适客舱，尊享空间
    无论是 B787 或 A330 的公务舱，还是明珠经济舱与经济舱，我们都为您提供宽敞舒适的座椅、先进的娱乐系统和贴心的服务。轻轻一按，豪华座椅即可180°平躺，助您在万米高空安享睡眠。
    
    🍱 精致餐食，品味云端美味
    南航为您准备了丰富多样的机上美食：从地道粤菜到国际风味，从岭南特色煲仔饭到招牌南航大碗面，满足您的味蕾享受。更有绿色飞行服务，支持环保的同时还可获得里程奖励！
    
    🎧 海量娱乐，畅享空中时光
    南航天空影院精选百部热门影片，天空电视提供300+集电视剧与纪录片，上千首音乐作品任你点播，还有专为儿童打造的“儿童天地”频道，让全家人都能尽享旅途乐趣！
    
    🌐 高速上网，云端互联不停歇
    部分航班提供全程高速 Wi-Fi 上网服务，让您在飞行途中也能轻松处理工作、浏览资讯或与亲友保持联系。
    
    🛍️ 免税精品，全球好物随心购
    南航机上免税品涵盖香水、护肤品、彩妆、名表、珠宝、烟酒等高端商品，价格优惠，品质保障。还可以提前预订，下机前直接领取，方便又省心！
    
    🌱 绿色飞行，关爱地球未来
    参与“绿色飞行”计划，按需用餐，减少浪费，还能获得里程奖励或保险福利，为环保出一份力！
    
    📸 个性服务，贴心体验
    特殊餐食预订、婴儿餐、宗教餐、医疗餐等多样化需求我们全力满足；还有专业的乘务团队为您提供细致入微的服务。
    
    —— 选南航，不只是飞得远，更是飞得舒适、飞得精彩！
    
    如您想了解更多关于南航服务详情，欢迎随时向我提问，我将为您详细介绍每一项贴心服务（机上餐食、机舱特色、机上娱乐、机上上网、机上免税品服务）！
    ###
    """
    bot = Assistant(llm=llm_cfg,
                    function_list=tools,
                    name='航班推荐助手',
                    system_message=system_message,
                    description="我是航班查询助手，可根据您的需求为您推荐航班。")

    return bot

def init_agent_service_amaps():
    llm_cfg = {
        # Use the model service provided by DashScope:
        'model': 'qwen3-235b-a22b',
        'model_server': 'dashscope',
        'api_key': "sk-1eb585dbfeea444691cd4216fd2a071d",
        'generate_cfg': {
            'top_p': 0.8
        }
    }

    tools = [
        {
            "mcpServers": {
                "AISearch": {
                    "url": "http://appbuilder.baidu.com/v2/ai_search/mcp/sse?api_key=Bearer+bce-v3/ALTAK-zHZ0uMemmKiBJD6n3o4UX/9d081c95580f5b9d37ffacb63f815f4c441e8fc3"
                }
            }
        },
        # 高德地图MCP server
        {
            "mcpServers": {
                "amap-maps": {
                    "command": "npx",
                    "args": [
                        "-y",
                        "@amap/amap-maps-mcp-server"
                    ],
                    "env": {
                        "AMAP_MAPS_API_KEY": "5b4f21c9dc29219b19cb7666258245c1"
                    }
                }
            }
        },
        'get_current_datetime',
        'code_interpreter',  # Built-in tools
    ]

    system_message = """
    你是一个专业的旅行规划助手，请注意现在是2025年。
    以下是默认操作：
    1) 如果需要为用户生成旅行规划，请调用amap-maps的高德地图MCP Server工具查询景点，进行旅行规划，保证规划的实际性。 
    2) 请一定要在回答的的尾部显式地给出信息来源，显式给出信息是来自于你固有知识还是外部信息，如（景点信息来源：高德地图）。请注意，严禁编造虚假的信息来源，如“途牛旅行”等为附加给智能体的信息源。
    
    以下是非默认操作，仅在用户要求下进行,如果用户不要求则不进行以下操作：
    1) 如果用户需要图文并茂的旅游攻略，请你通过amap-maps的amap-maps-maps_text_search工具获取图像的url，最好可以做到攻略中推荐的每个景点都跟随一张图像，多模态呈现。
    """

    bot = Assistant(llm=llm_cfg,
                    function_list=tools,
                    name='旅行规划助手',
                    system_message=system_message,
                    description="我是一个基于高德地图的旅行规划助手，可根据您的喜好为您定制行程。")

    return bot

def init_agent_service_redbook():
    llm_cfg = {
        # Use the model service provided by DashScope:
        'model': 'qwen3-235b-a22b',
        'model_server': 'dashscope',
        'api_key': "sk-1eb585dbfeea444691cd4216fd2a071d",
        'generate_cfg': {
            'top_p': 0.8
        }
    }

    tools = [
        {
            "mcpServers": {
                "AISearch": {
                    "url": "http://appbuilder.baidu.com/v2/ai_search/mcp/sse?api_key=Bearer+bce-v3/ALTAK-zHZ0uMemmKiBJD6n3o4UX/9d081c95580f5b9d37ffacb63f815f4c441e8fc3"
                }
            }
        },
        # 小红书
        {
            "mcpServers": {
                "xiaohongshu MCP": {
                    "command": current_interpreter,  # 配置为解释器地址
                    "args": [
                        os.path.join(BASE_DIR, "Redbook-Search-Comment-MCP2.0", "xiaohongshu_mcp.py"),
                        "--stdio"
                    ]
                }
            }
        },
        'get_current_datetime',
        'code_interpreter',  # Built-in tools
    ]

    system_message = """
    你是一个基于浏览器和小红书的专业的旅行规划助手，请注意现在是2025年。
    以下是默认操作：
    1) 请一定要在回答的的尾部给出信息来源，显式给出信息是来自于你固有知识还是外部信息，如（高铁信息来源：小红书）。
    2) 如果需要为用户生成旅行规划，请根据小红书图文生成旅行攻略。 
    """

    bot = Assistant(llm=llm_cfg,
                    function_list=tools,
                    name='小红书出行助手',
                    system_message=system_message,
                    description="我是一个基于小红书旅行规划助手，可根据小红书总结旅行攻略。")

    return bot

def init_agent_service_webpage():
    llm_cfg = {
        # Use the model service provided by DashScope:
        'model': 'qwen3-235b-a22b',
        'model_server': 'dashscope',
        'api_key': "sk-1eb585dbfeea444691cd4216fd2a071d",
        'generate_cfg': {
            'top_p': 0.8
        }
    }

    tools = [
        # Edgeone page生成网页
        {
            "mcpServers": {
                "edgeone-pages-mcp-server": {
                    "command": "npx",
                    "args": ["edgeone-pages-mcp"],
                    "env": {
                        "EDGEONE_PAGES_API_TOKEN": "VqQ91Y9EMJtfK8pvPQzS+kPIHsBi8C9bb4eYi5If9+4=",
                        "EDGEONE_PAGES_PROJECT_NAME": os.path.join(BASE_DIR)  # 配置为现有项目地址
                    }
                }
            }
        },
        'get_current_datetime',
        'code_interpreter',  # Built-in tools
    ]

    system_message = """
    你是一个网页生成助手。请将用户确定的旅行攻略包装成美观的网页。
    """

    bot = Assistant(llm=llm_cfg,
                    function_list=tools,
                    name='网页生成助手',
                    system_message=system_message,
                    description="我是一个网页包装助手。")

    return bot

def app_gui():
    # Define the agent
    bot_railway_plan = init_agent_service_railway_plan()
    bot_flight_plan = init_agent_service_flight_plan()
    bot_amaps = init_agent_service_amaps()
    bot_redbook = init_agent_service_redbook()
    bot_webpage = init_agent_service_webpage()
    bot_list = [bot_railway_plan, bot_flight_plan, bot_amaps, bot_redbook, bot_webpage]

    chatbot_config = {
        'prompt.suggestions': [
            '请帮我规划南京三天两晚的旅行计划。',
            '请帮我查询2025年6月12日从广州前往杭州的航班。',
            '请推荐明天上午从杭州前往上海的高铁。',
            '帮我找三篇南京旅游的小红书图文，并总结成三天两晚的旅行攻略，附上图文链接。'
        ]
    }
    WebUI(bot_list, chatbot_config=chatbot_config).run(enable_mention=True, share=True)

if __name__ == '__main__':

    app_gui()