# Copyright 2023 The Qwen team, Alibaba Group. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A sqlite database assistant implemented by assistant"""

import argparse
import asyncio
import json
import os
from typing import Dict, List, Optional, Union

from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
from qwen_agent.llm.oai import TextChatAtOAI
from qwen_agent.llm.schema import ASSISTANT, FUNCTION, Message
from qwen_agent.utils.output_beautify import typewriter_print

ROOT_RESOURCE = os.path.join(os.path.dirname(__file__), "resource")




def read_mcp_settings():
    with open("./mcp/mcp_settings.json", "r") as f:
        mcp_settings_json = json.load(f)
    return mcp_settings_json


def read_mcp_settings_with_project_restriction(project_data_dir: str):
    """读取MCP配置并添加项目目录限制"""
    with open("./mcp/mcp_settings.json", "r") as f:
        mcp_settings_json = json.load(f)
    
    # 为json-reader添加项目目录限制
    for server_config in mcp_settings_json:
        if "mcpServers" in server_config:
            for server_name, server_info in server_config["mcpServers"].items():
                if server_name == "json-reader":
                    # 添加环境变量来传递项目目录限制
                    if "env" not in server_info:
                        server_info["env"] = {}
                    server_info["env"]["PROJECT_DATA_DIR"] = project_data_dir
                    server_info["env"]["PROJECT_ID"] = project_data_dir.split("/")[-2] if "/" in project_data_dir else "default"
                    break
    
    return mcp_settings_json


def read_system_prompt():
    with open("./agent_prompt.txt", "r", encoding="utf-8") as f:
        return f.read().strip()


def read_system_prompt():
    """读取通用的无状态系统prompt"""
    with open("./agent_prompt.txt", "r", encoding="utf-8") as f:
        return f.read().strip()


def init_agent_service():
    """默认初始化函数，保持向后兼容"""
    return init_agent_service_universal("qwen3-next")


def read_llm_config():
    """读取LLM配置文件"""
    with open("./llm_config.json", "r") as f:
        return json.load(f)


def init_agent_service_with_project(project_id: str, project_data_dir: str, model_name: str = "qwen3-next"):
    """支持项目目录的agent初始化函数 - 保持向后兼容"""
    llm_cfg = read_llm_config()
      
    # 读取通用的系统prompt（无状态）
    system = read_system_prompt()

    # 读取MCP工具配置
    tools = read_mcp_settings_with_project_restriction(project_data_dir)
    
    # 使用指定的模型配置
    if model_name not in llm_cfg:
        raise ValueError(f"Model '{model_name}' not found in llm_config.json. Available models: {list(llm_cfg.keys())}")
    
    llm_instance = llm_cfg[model_name]
    if "llm_class" in llm_instance:
        llm_instance = llm_instance.get("llm_class", TextChatAtOAI)(llm_instance)

    bot = Assistant(
        llm=llm_instance,  # 使用指定的模型实例
        name=f"数据库助手-{project_id}",
        description=f"项目 {project_id} 数据库查询",
        system_message=system,
        function_list=tools,
    )

    return bot


def init_agent_service_universal():
    """创建无状态的通用助手实例（使用默认LLM，可动态切换）"""
    llm_cfg = read_llm_config()
    
    # 读取通用的系统prompt（无状态）
    system = read_system_prompt()
    
    # 读取基础的MCP工具配置（不包含项目限制）
    tools = read_mcp_settings()

    # 使用默认模型创建助手实例
    default_model = "qwen3-next"  # 默认模型
    if default_model not in llm_cfg:
        # 如果默认模型不存在，使用第一个可用模型
        default_model = list(llm_cfg.keys())[0]
    
    llm_instance = llm_cfg[default_model]
    if "llm_class" in llm_instance:
        llm_instance = llm_instance.get("llm_class", TextChatAtOAI)(llm_instance)

    bot = Assistant(
        llm=llm_instance,  # 使用默认LLM初始化
        name="通用数据检索助手",
        description="无状态通用数据检索助手",
        system_message=system,
        function_list=tools,
    )

    return bot


def update_agent_llm(agent, model_name: str, api_key: str = None, model_server: str = None):
    """动态更新助手实例的LLM，支持从接口传入参数"""

    # 获取基础配置
    llm_config =  {
        "model": model_name,
        "api_key": api_key,
        "model_server": model_server
    }
    
    # 创建LLM实例，确保不是字典
    #if "llm_class" in llm_config:
    #    llm_instance = llm_config.get("llm_class", TextChatAtOAI)(llm_config)
    #else:
    # 使用默认的 TextChatAtOAI 类
    llm_instance = TextChatAtOAI(llm_config)
    
    # 动态设置LLM
    agent.llm = llm_instance
    
    return agent


def test(query="数据库里有几张表"):
    # Define the agent
    bot = init_agent_service()

    # Chat
    messages = []

    messages.append({"role": "user", "content": query})

    responses = []
    for response in bot.run(messages):
        responses.append(response)

    # 只输出最终结果，不显示中间过程
    if responses:
        final_response = responses[-1][-1]  # 取最后一个响应作为最终结果
        print("Answer:", final_response["content"])


def app_tui():
    # Define the agent
    bot = init_agent_service()

    # Chat
    messages = []
    while True:
        # Query example: 数据库里有几张表
        query = input("user question: ")
        # File example: resource/poem.pdf
        file = input("file url (press enter if no file): ").strip()
        if not query:
            print("user question cannot be empty！")
            continue
        if not file:
            messages.append({"role": "user", "content": query})
        else:
            messages.append(
                {"role": "user", "content": [{"text": query}, {"file": file}]}
            )

        response = []
        for response in bot.run(messages):
            print("bot response:", response)
        messages.extend(response)


def app_gui():
    # Define the agent
    bot = init_agent_service()
    chatbot_config = {
        "prompt.suggestions": [
            "数据库里有几张表",
            "创建一个学生表包括学生的姓名、年龄",
            "增加一个学生名字叫韩梅梅，今年6岁",
        ]
    }
    WebUI(
        bot,
        chatbot_config=chatbot_config,
    ).run()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="数据库助手")
    parser.add_argument(
        "--query", type=str, default="数据库里有几张表", help="用户问题"
    )
    parser.add_argument(
        "--mode",
        type=str,
        choices=["test", "tui", "gui"],
        default="test",
        help="运行模式",
    )
    args = parser.parse_args()

    if args.mode == "test":
        test(args.query)
    elif args.mode == "tui":
        app_tui()
    elif args.mode == "gui":
        app_gui()
