import os

from llama_index.core.settings import Settings
from llama_index.core.agent import AgentRunner, StructuredPlannerAgent, FunctionCallingAgentWorker, ReActChatFormatter
from llama_index.core.tools.query_engine import QueryEngineTool
from sqlalchemy import create_engine, Engine
from llama_index.core.response_synthesizers.type import ResponseMode

from app.engine.engine import create_query_engine, create_summary_query_engine
from app.engine.loaders.db import makeDescriptionByEngine
from app.engine.tools import ToolFactory
from app.engine.index import get_index


def get_chat_engine(filters=None, params=None):
    system_prompt = os.getenv("SYSTEM_PROMPT")
    top_k = int(os.getenv("TOP_K", "3"))
    tools = []

    # 创建SQL查询工具
    sql_query_engine = create_summary_query_engine()
    sql_query_tool = QueryEngineTool.from_defaults(query_engine=sql_query_engine,
                                                   name="zjdata_query_tool",
                                                   description="来源于一个由博微公司电力造价软件编制的造价工程文件。该文件以多张表格的形式存储存储了整个工程的全部数据内容。适用于以详细的自然语言查询表格数据方式查询造价工程各项具体属性、费用的数值。请先使用“zj_query_tool”无法解决才使用本工具"
                                                   )
    #tools.append(sql_query_tool)

    # Add query tool if index exists
    index = get_index()
    if index is not None:
        summary_query_engine = create_summary_query_engine(index)
        summary_query_tool = QueryEngineTool.from_defaults( query_engine=summary_query_engine, name="summary_query_tool",
                                                            description="适用于任何需要进行全面总结、概括的要求。",
                                                            )
        query_engine = create_query_engine()
        query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine, name="zj_query_tool",
                                                          )

        tools.append(summary_query_tool)
        tools.append(query_engine_tool)

    # Add additional tools
    tools += ToolFactory.from_env()

    prefix_messages = ("""您的设计旨在帮助完成各种任务，从回答问题到提供其他类型分析的摘要。\n\n##工具\n\n你可以访问各种工具。你有责任按照你认为合适的顺序使用这些工具来完成当前的任务。\n这可能需要将任务分解为子任务，并使用不同的工具来完成每个子任务。\n\n你可以访问以下工具：\n{tool_desc}\n\n\n##输出格式\n\n请用与问题相同的语言回答，并使用以下格式：\n\n```\nThought: 用户当前的语言是：(user's language)。我需要使用工具来帮助我回答问题。\nAction: 如果使用工具,则为工具名称(one of {tool_names})。\nAction Input: 输入给工具的内容，使用JSON格式表示kwargs（例如{{\"input\": \"hello world\", \"num_beams\": 5}}）\n```\n\n请始终以Thought开始。\n\n切勿用Markdown代码标记包围你的响应。如果需要，可以在响应中使用代码标记。\n\n请为Action Input使用有效的JSON格式。不要这样做{{\'input\': \'hello world\', \'num_beams\': 5}}。\n\n如果使用此格式，用户将以下面的格式进行回应：\n\n```\nObservation: 工具响应\n```\n\n你应该继续重复上述格式，直到你有足够的信息来回答问题而无需使用更多工具。此时，你必须使用以下两种格式之一进行回答：\n\n```\nThought: 我可以不用任何工具来回答。我将使用用户的语言来回答。\nAnswer: [你的答案（与用户问题相同的语言）]\n```\n\n```\nThought: 我无法使用提供的工具回答问题。\nAnswer: [你的答案（与用户问题相同的语言）]\n```\n\n## 当前对话\n\n以下是当前对话，由人类和助手的消息交替组成。\n""")
    react_chat_formatter = ReActChatFormatter.from_defaults(prefix_messages)
    agentrunner = AgentRunner.from_llm(
        llm=Settings.llm,
        tools=tools,
        react_chat_formatter=react_chat_formatter,
        system_prompt=system_prompt,
        verbose=True,
    )
    return agentrunner
    # create the function calling worker for reasoning
    # worker = FunctionCallingAgentWorker.from_tools(
    #     tools, verbose=True
    # )
    #
    # # wrap the worker in the top-level planner
    # return StructuredPlannerAgent(worker, tools)