import os

from llama_index.core import SQLDatabase, SummaryIndex, VectorStoreIndex
from llama_index.core.indices.struct_store import SQLTableRetrieverQueryEngine
from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex
from llama_index.core.settings import Settings
from llama_index.core.agent import AgentRunner, StructuredPlannerAgent, FunctionCallingAgentWorker
from llama_index.core.tools.query_engine import QueryEngineTool
from sqlalchemy import create_engine, Engine

from app.engine.loaders.db import makeDescriptionByEngine
from app.engine.tools import ToolFactory
from app.engine.index import get_index
from app.settings import get_node_postprocessors

sql_database = None
sql_obj_index = None

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

    global sql_obj_index
    global sql_database
    if sql_obj_index is None:
        sqlengine = create_engine(os.getenv("SQL_DATABASE_URL", ""))
        sql_database = SQLDatabase(sqlengine)
        table_schema_objs = makeDescriptionByEngine(sql_database)
        table_node_mapping = SQLTableNodeMapping(sql_database)

        sql_obj_index = ObjectIndex.from_objects(
            table_schema_objs,
            table_node_mapping,
            index_cls=VectorStoreIndex,
        )

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

    # Add query tool if index exists
    index = get_index()
    if index is not None:
        summary_index = SummaryIndex(index.vector_store.get_nodes(node_ids=None))
        summary_query_engine = summary_index.as_query_engine()
        summary_query_tool = QueryEngineTool.from_defaults( query_engine=summary_query_engine, name="summary_query_tool",
                                                            description="适用于任何需要进行全面总结、概括的要求。",
                                                            #description="适用于任何需要对所有内容进行全面总结的请求。有关电力造价领域更具体部分的问题，请使用zj_query_engine_tool",
                                                            )

        # 创建向量检索查询工具
        postprocess = get_node_postprocessors()
        query_engine = index.as_query_engine(
            similarity_top_k=top_k, filters=filters,
            node_postprocessors=postprocess,
        )
        query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine, name="zj_query_tool",
                                                          description="由博微公司编制的关于电力造价知识、电力造价编制软件知识和造价工程文件结构的知识库。适用于查询电力领域、电力造价领域、博微、博微电力、博微造价等业务等内容。如果本知识库没有直接答案但有解决思路的可以返回解决办法后建议使用“zjdata_query_tool”工具。如果你不知道答案，就说你不知道，不要编造答案。",
                                                          )

        tools.append(summary_query_tool)
        tools.append(query_engine_tool)
        #tools.append(sql_query_tool)

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

    return AgentRunner.from_llm(
        llm=Settings.llm,
        tools=tools,
        system_prompt=system_prompt,
        verbose=True,
    )
    # 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)